From 59d399da7d8e2d653b0f6d9f7bb135eee6776fb7 Mon Sep 17 00:00:00 2001
From: xiaozhon Huang <xiaozhong.huang@cern.ch>
Date: Mon, 9 Nov 2020 17:50:40 +0000
Subject: [PATCH] tauRecTools: associate cluster in CaloCalTopoClusterContainer
 to the tau candidate

For Topo seed jets, the cluster links in the tau candidate point to the
clusters in a shallow copy container, which will not be available in AOD
of R22. To fix this, the cluster links point to the cluster in the
original container.

For PFlow seed jets, the cluster links in the tau candidate point to the
PFOs(neutral or charged). It is confusing. Now the cluster links point
to the corresponding cluter in the PFO object.

Previously, only clusters within 0.2 cone of tau candidate is stored. In
the tau reconstruction, we need all the clusters. Thus the cut of dR is now
removed.

The clusters could have energy <=0, which will be skimmed in AOD. The
link then becomes invalid. Since we do not want an invalid link, all
the clusters are requied to have positive energy.
---
 Event/xAOD/xAODTau/Root/TauJet_v3.cxx         |  46 +++++++-
 .../xAOD/xAODTau/xAODTau/versions/TauJet_v3.h |   4 +
 .../Root/METTauAssociator.cxx                 |  22 +---
 .../tauRec/python/TauAlgorithmsHolder.py      |  15 +++
 Reconstruction/tauRec/python/TauRecBuilder.py |   2 +
 .../tauRecTools/src/TauAxisSetter.cxx         |   7 --
 .../tauRecTools/src/TauClusterFinder.cxx      | 105 ++++++++++++++++++
 .../tauRecTools/src/TauClusterFinder.h        |  42 +++++++
 .../src/components/tauRecTools_entries.cxx    |   2 +
 9 files changed, 219 insertions(+), 26 deletions(-)
 create mode 100644 Reconstruction/tauRecTools/src/TauClusterFinder.cxx
 create mode 100644 Reconstruction/tauRecTools/src/TauClusterFinder.h

diff --git a/Event/xAOD/xAODTau/Root/TauJet_v3.cxx b/Event/xAOD/xAODTau/Root/TauJet_v3.cxx
index 2f62bb0434d..7f5670867d9 100644
--- a/Event/xAOD/xAODTau/Root/TauJet_v3.cxx
+++ b/Event/xAOD/xAODTau/Root/TauJet_v3.cxx
@@ -17,7 +17,7 @@
 #include "xAODTau/versions/TauJetCalibMapper_v1.h"
 //#include "xAODTau/versions/TauJetCalibMapper_v3.h"
 #include "TauJetAccessors_v3.h"
-
+#include "xAODCaloEvent/CaloVertexedTopoCluster.h"
 
 namespace{
   bool
@@ -627,6 +627,50 @@ namespace xAOD {
     return *(clusterAcc(*this).at(i));
   }
 
+
+  std::vector<const IParticle*> TauJet_v3::clusters(double dRMax) const {
+    std::vector<const IParticle*> particleList;
+
+    for (const auto& link : clusterAcc(*this)) {
+      const IParticle* particle = *link;
+
+      if (dRMax > 0.0) {
+        const xAOD::CaloCluster* cluster = static_cast<const xAOD::CaloCluster*>(particle);
+        const xAOD::Vertex* vertex = nullptr;
+        TLorentzVector clusterP4;
+        TLorentzVector tauP4;
+        
+        // Apply the vertex correction if there is a vertex
+        if (this->vertexLink().isValid()) {
+          vertex = this->vertex();
+          tauP4 = this->p4(xAOD::TauJetParameters::IntermediateAxis);  
+        }
+        else {
+          const xAOD::Jet* jet = this->jet();
+          bool isAvailable = jet->getAssociatedObject("OriginVertex", vertex);
+          if (!isAvailable) vertex = nullptr;
+          tauP4 = this->p4(xAOD::TauJetParameters::DetectorAxis);
+        }
+
+        if (vertex) {
+          clusterP4 = xAOD::CaloVertexedTopoCluster(*cluster, vertex->position()).p4();
+        }
+        else {
+          clusterP4 = cluster->p4();
+        }
+
+        double dR = clusterP4.DeltaR(tauP4);
+        if (dR > dRMax) continue;
+      }
+
+      particleList.push_back(particle);
+    }
+
+    return particleList;
+  }
+
+
+
   TauJet_v3::FourMom_t TauJet_v3::calibratedCluster( size_t i, xAOD::CaloCluster::State state/*=xAOD::CaloCluster::State::CALIBRATED*/) const{
     const xAOD::IParticle* part = this->cluster(i);
     if(!part) return FourMom_t();
diff --git a/Event/xAOD/xAODTau/xAODTau/versions/TauJet_v3.h b/Event/xAOD/xAODTau/xAODTau/versions/TauJet_v3.h
index 25141fbba29..48e52104518 100644
--- a/Event/xAOD/xAODTau/xAODTau/versions/TauJet_v3.h
+++ b/Event/xAOD/xAODTau/xAODTau/versions/TauJet_v3.h
@@ -291,6 +291,10 @@ namespace xAOD {
     void setClusterLinks( const IParticleLinks_t& clusters );
     /// Get the pointer to a given cluster associated with this tau
     const IParticle* cluster( size_t i) const;
+    
+    //* @brief Get the clusters within dR cone of the tau candidate */
+    std::vector<const IParticle*> clusters(double dR = 0.2) const;
+    
     /// Get TLV to a given cluster in calibrated state
     FourMom_t calibratedCluster( size_t i, xAOD::CaloCluster::State state=xAOD::CaloCluster::State::CALIBRATED) const;
     //number of cluster with associated to tau
diff --git a/Reconstruction/MET/METReconstruction/Root/METTauAssociator.cxx b/Reconstruction/MET/METReconstruction/Root/METTauAssociator.cxx
index cfafb090b37..d1f1316fd14 100644
--- a/Reconstruction/MET/METReconstruction/Root/METTauAssociator.cxx
+++ b/Reconstruction/MET/METReconstruction/Root/METTauAssociator.cxx
@@ -107,24 +107,10 @@ namespace met {
                                                    const met::METAssociator::ConstitHolder& /*tcCont*/) const
   {
     const TauJet* tau = static_cast<const TauJet*>(obj);
-    for( ElementLink< xAOD::IParticleContainer > cluster_link : tau->clusterLinks() ){
-      const xAOD::IParticle* ipart = *cluster_link;
-      if (ipart->type() == xAOD::Type::ParticleFlow){
-        // If using PFO, get cluster
-        const xAOD::PFO *pfo = static_cast<const xAOD::PFO*>(ipart);
-        if (pfo->isCharged()){ continue; }
-        else {
-          ipart = pfo->cluster(0);
-        }
-      }
-      if(ipart->type() != xAOD::Type::CaloCluster) {
-            ATH_MSG_WARNING("Unexpected jet constituent type " << ipart->type() << " received! Skip.");
-        continue;
-      }      
-      // Link set in Reconstruction/tauRecTools/src/TauAxisSetter.cxx
-      // Internal defaults are m_clusterCone = 0.2, m_doCellCorrection = false, m_doAxisCorrection = True
-      const CaloCluster* pClus = static_cast<const CaloCluster*>( ipart );
-      tclist.push_back(pClus);
+
+    for (const xAOD::IParticle* particle : tau->clusters(0.2)) {
+      const CaloCluster* cluster = static_cast<const CaloCluster*>(particle);
+      tclist.push_back(cluster);
     }
 
     return StatusCode::SUCCESS;
diff --git a/Reconstruction/tauRec/python/TauAlgorithmsHolder.py b/Reconstruction/tauRec/python/TauAlgorithmsHolder.py
index 749c9a83751..bd53f55586d 100644
--- a/Reconstruction/tauRec/python/TauAlgorithmsHolder.py
+++ b/Reconstruction/tauRec/python/TauAlgorithmsHolder.py
@@ -667,6 +667,21 @@ def getTauTrackFinder(removeDuplicateTracks=True):
     cached_instances[_name] = TauTrackFinder      
     return TauTrackFinder
 
+
+# Associate the cluster in jet constituents to the tau candidate
+def getTauClusterFinder():
+    _name = sPrefix + 'TauClusterFinder'
+
+    if _name in cached_instances:
+        return cached_instances[_name]
+
+    from tauRecTools.tauRecToolsConf import TauClusterFinder
+    TauClusterFinder = TauClusterFinder(name = _name)
+
+    cached_instances[_name] = TauClusterFinder
+    return TauClusterFinder
+
+
 ########################################################################
 # MvaTESVariableDecorator
 def getMvaTESVariableDecorator():
diff --git a/Reconstruction/tauRec/python/TauRecBuilder.py b/Reconstruction/tauRec/python/TauRecBuilder.py
index e97f81ba5cb..24bf9548f7d 100644
--- a/Reconstruction/tauRec/python/TauRecBuilder.py
+++ b/Reconstruction/tauRec/python/TauRecBuilder.py
@@ -58,6 +58,8 @@ class TauRecCoreBuilder ( TauRecConfigured ) :
                 tools.append(taualgs.getTauVertexFinder(doUseTJVA=self.do_TJVA))
             tools.append(taualgs.getTauAxis())
             tools.append(taualgs.getTauTrackFinder(removeDuplicateTracks=(tauFlags.removeDuplicateCoreTracks() ) ))
+            tools.append(taualgs.getTauClusterFinder())
+
             if doMVATrackClassification : tools.append(taualgs.getTauTrackClassifier())
             if not doMVATrackClassification and doRNNTrackClassification:
                 tools.append(taualgs.getTauTrackRNNClassifier())
diff --git a/Reconstruction/tauRecTools/src/TauAxisSetter.cxx b/Reconstruction/tauRecTools/src/TauAxisSetter.cxx
index ad87dbd4f4f..cb465266aa7 100644
--- a/Reconstruction/tauRecTools/src/TauAxisSetter.cxx
+++ b/Reconstruction/tauRecTools/src/TauAxisSetter.cxx
@@ -57,13 +57,6 @@ StatusCode TauAxisSetter::execute(xAOD::TauJet& pTau) const {
     double dR = baryCenter.DeltaR(constituentP4);
     if (dR > m_clusterCone) continue;
 
-    ElementLink<xAOD::IParticleContainer> linkToCluster;
-    linkToCluster.toContainedElement(
-      *(static_cast<const xAOD::IParticleContainer*> (constituent->rawConstituent()->container())), 
-      constituent->rawConstituent()
-      );
-    pTau.addClusterLink(linkToCluster);
-
     tauDetectorAxis += constituentP4;
     ++nConstituents;
   }
diff --git a/Reconstruction/tauRecTools/src/TauClusterFinder.cxx b/Reconstruction/tauRecTools/src/TauClusterFinder.cxx
new file mode 100644
index 00000000000..8fbe9885ad3
--- /dev/null
+++ b/Reconstruction/tauRecTools/src/TauClusterFinder.cxx
@@ -0,0 +1,105 @@
+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef XAOD_ATHLYSIS
+
+#include "TauClusterFinder.h"
+#include "tauRecTools/HelperFunctions.h"
+
+#include "xAODTau/TauJetContainer.h"
+#include "xAODTau/TauJetAuxContainer.h"
+#include "xAODTau/TauJet.h"
+#include "xAODBase/IParticleHelpers.h"
+
+
+TauClusterFinder::TauClusterFinder(const std::string& name) :
+TauRecToolBase(name) {
+}
+
+
+
+StatusCode TauClusterFinder::execute(xAOD::TauJet& tau) const {
+  if (! tau.jetLink().isValid()) {
+    ATH_MSG_ERROR("Tau jet link is invalid.");
+    return StatusCode::FAILURE;
+  }
+  const xAOD::Jet* jetSeed = tau.jet();
+  
+  // Find all the clusters in the JetConstituent
+  std::vector<const xAOD::CaloCluster*> clusterList = getClusterList(*jetSeed);
+
+  for (const xAOD::CaloCluster* cluster : clusterList) {
+    // Clusters with negative energy will be thinned, and the elementlinks to these
+    // clusters will not be valid. 
+    if (cluster->e() <=0) continue;
+
+    ElementLink<xAOD::IParticleContainer> linkToCluster;
+    linkToCluster.toContainedElement(
+      *(static_cast<const xAOD::IParticleContainer*> (cluster->container())), 
+      static_cast<const xAOD::IParticle*>(cluster));
+    tau.addClusterLink(linkToCluster);
+  }
+
+  return StatusCode::SUCCESS;
+}
+
+
+
+std::vector<const xAOD::CaloCluster*> TauClusterFinder::getClusterList(const xAOD::Jet& jet) const {
+  std::vector<const xAOD::CaloCluster*> clusterList;
+
+  xAOD::JetConstituentVector constituents = jet.getConstituents();
+  for (const xAOD::JetConstituent* constituent : constituents) {
+    ATH_MSG_DEBUG("JetConstituent: ");
+    ATH_MSG_DEBUG("eta: " << constituent->eta() << " phi: " << constituent->phi() << " e: " << constituent->e()); 
+
+    if( constituent->type() == xAOD::Type::CaloCluster ) {
+	  const xAOD::CaloCluster* cluster = static_cast<const xAOD::CaloCluster*>(constituent->rawConstituent());
+      ATH_MSG_DEBUG("CaloCluster: ");
+	  ATH_MSG_DEBUG("eta: " << cluster->eta() << " phi: " << cluster->phi() << " e: " << cluster->e());
+	  
+      // Cluster in Topo jets is a shallow copy of CaloCalTopoCluster. Since jets may not be stored to AOD in R22,
+      // we will not be able retrieve the cluster from the jets with AOD as input. To fix this, we retrieve
+      // the cluster in the orignial container, i.e. CaloCalTopoCluster.
+      const xAOD::IParticle* originalParticle = xAOD::getOriginalObject(*cluster);
+      const xAOD::CaloCluster* orignialCluster = static_cast<const xAOD::CaloCluster*>(originalParticle);
+	  clusterList.push_back(orignialCluster);
+    }
+    else if ( constituent->type() == xAOD::Type::ParticleFlow ) {
+	  const xAOD::PFO* pfo = static_cast<const xAOD::PFO*>(constituent->rawConstituent());
+	  
+      if (pfo->isCharged()) continue;
+	  if (pfo->nCaloCluster()!=1){
+	    ATH_MSG_WARNING("Neutral PFO has " << std::to_string(pfo->nCaloCluster()) << " clusters, expected exactly 1!\n");
+        continue;
+      }
+	  
+      clusterList.push_back(pfo->cluster(0));
+    }
+    else {
+	  ATH_MSG_WARNING("Seed jet constituent type not supported ! Skip");
+	  continue;
+    }
+  }
+  
+  for (const xAOD::JetConstituent* constituent : constituents) {
+    // There is only one type in the constituents
+    if (constituent->type() != xAOD::Type::ParticleFlow) break;
+	
+    const xAOD::PFO* pfo = static_cast<const xAOD::PFO*>( constituent->rawConstituent() );
+	if (! pfo->isCharged()) continue;
+
+	for (u_int index=0; index<pfo->nCaloCluster(); index++){
+	  const xAOD::CaloCluster* cluster = pfo->cluster(index);
+	  // Check it is not duplicate of one in neutral list
+	  if ( std::find(clusterList.begin(), clusterList.end(), cluster) == clusterList.end() ) {
+	    clusterList.push_back(cluster);
+	  }
+	}
+  }
+ 
+  return clusterList;
+} 
+
+#endif
diff --git a/Reconstruction/tauRecTools/src/TauClusterFinder.h b/Reconstruction/tauRecTools/src/TauClusterFinder.h
new file mode 100644
index 00000000000..10ade8d4ed5
--- /dev/null
+++ b/Reconstruction/tauRecTools/src/TauClusterFinder.h
@@ -0,0 +1,42 @@
+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
+*/
+
+#ifndef TAUCLUSTERFINDER_H
+#define TAUCLUSTERFINDER_H
+
+#include "tauRecTools/TauRecToolBase.h"
+#include "tauRecTools/ITauVertexCorrection.h"
+
+#include "AsgTools/ToolHandle.h"
+
+/**
+ * @brief Associate the clusters used in the seed jet to the tau candidate. 
+ * 
+ *  The asscated clusters will always be the clusters in CaloCalTopoClusterContainer. 
+ * 
+ * @author Xiaozhong Huang <xiaozhong.huang@cern.ch>
+ *                                                                              
+ */
+
+class TauClusterFinder : public TauRecToolBase {
+
+  public:
+    
+    ASG_TOOL_CLASS2(TauClusterFinder, TauRecToolBase, ITauToolBase);
+    
+    /** @brief Constructor */ 
+    TauClusterFinder(const std::string& name);
+
+    /** @brief Destructor */
+    ~TauClusterFinder() = default;
+
+    /** @brief Execution of this tool */ 
+    virtual StatusCode execute(xAOD::TauJet& tau) const override;
+
+  private:
+
+    std::vector<const xAOD::CaloCluster*> getClusterList(const xAOD::Jet& jet) const;
+};
+
+#endif
diff --git a/Reconstruction/tauRecTools/src/components/tauRecTools_entries.cxx b/Reconstruction/tauRecTools/src/components/tauRecTools_entries.cxx
index 476f9492243..4bee4c3b26a 100644
--- a/Reconstruction/tauRecTools/src/components/tauRecTools_entries.cxx
+++ b/Reconstruction/tauRecTools/src/components/tauRecTools_entries.cxx
@@ -3,6 +3,7 @@
 #include "../TauAxisSetter.h"
 #include "../TauCellVariables.h"
 #include "../TauTrackFinder.h"
+#include "../TauClusterFinder.h"
 #include "../TauVertexFinder.h"
 #include "../TauElectronVetoVariables.h"
 #include "../TauShotFinder.h"
@@ -34,6 +35,7 @@ DECLARE_COMPONENT( JetSeedBuilder )
 DECLARE_COMPONENT( TauAxisSetter )
 DECLARE_COMPONENT( TauCellVariables )
 DECLARE_COMPONENT( TauTrackFinder )
+DECLARE_COMPONENT( TauClusterFinder )
 DECLARE_COMPONENT( TauVertexFinder )
 DECLARE_COMPONENT( TauElectronVetoVariables )
 DECLARE_COMPONENT( TauShotFinder )
-- 
GitLab