tauRec: improved tau cell thinning

Hello,

This MR is implementing an improved version of the tau pi0 cell thinning in xAODs (ATLASRECTS-5684, ATLTAU-1733).
The pi0 clusters are initially built from cells stored in a separate container (TauCommonPi0Cells).
An algorithm (ClusterCellRelinkAlg) is introduced to recreate pi0 clusters with cells from the main cell container (AllCalo), thereby eliminating the overhead from our separate cell container.
The ESD and AOD contents have been adjusted accordingly: TauCommonPi0Cells is dropped, and the relinked pi0s are written out instead of the initial pi0s.
There will be a follow-up MR that will drop some other pi0 container we used to use in TauPi0ClusterCreator, which should no longer be needed.
For now, only the minimum amount of TauPi0ClusterCreator code was changed.

Another algorithm (TauCellThinningAlg) is introduced to save the cells from the tau seed-jet topoclusters in xAODs.
Only clusters within 0.2 of the tau are considered, to limit the xAOD size increase.
Note: in the xAODs, we are still missing CaloCalTopoClusters_links (the CaloClusterCellLinkContainer) for the tau seed-jet clusters.
An attempt at thinning this container was made, but it was not successful.

Cheers,
Bertrand

Calorimeter/CaloExample/CaloRecEx/share/CaloRecOutputItemList_jobOptions.py

@@ -205,8 +205,6 @@ for theKey in CaloClusterKeys: #Fixme .. Apply this only to TopoClusters?
 
 CaloAODList+=CaloClusterItemList
 
-CaloAODList+=["CaloClusterContainer#Tau1P3PPi0ClusterContainer"]
-
 # E4' cells
 CaloAODList+=["TileCellContainer#E4prContainer"]
 

Reconstruction/RecExample/RecExCommon/share/RecExCommon_topOptions.py

@@ -1281,14 +1281,22 @@ if ( rec.doAOD() or rec.doWriteAOD()) and not rec.readAOD() :
                 addClusterToCaloCellAOD("InDetTrackParticlesAssociatedClusters")
 
             from tauRec.tauRecFlags import tauFlags
-            if ( rec.readESD() or tauFlags.Enabled() ) and rec.doTau:
+            if ( rec.readESD() or tauFlags.Enabled() ) and rec.doTau:                
                 from CaloRec.CaloRecConf import CaloThinCellsByClusterAlg
-                alg = CaloThinCellsByClusterAlg('CaloThinCellsByClusterAlg_TauInitialPi0Clusters',
+                alg = CaloThinCellsByClusterAlg('CaloThinCellsByClusterAlg_TauPi0Clusters',
                                                 StreamName = 'StreamAOD',
-                                                Clusters = 'TauInitialPi0Clusters',
-                                                Cells = 'TauCommonPi0Cells')
+                                                Clusters = 'TauPi0Clusters',
+                                                Cells = 'AllCalo')
                 topSequence += alg
 
+                from tauRec.tauRecConf import TauCellThinningAlg
+                alg = TauCellThinningAlg('TauCellThinningAlg',
+                                         StreamName = 'StreamAOD',
+                                         Cells = 'AllCalo',
+                                         Taus = "TauJets",
+                                         UseSubtractedCluster = tauFlags.useSubtractedCluster())
+                topSequence += alg
+                
         except Exception:
             treatException("Could not make AOD cells" )
 

Reconstruction/tauRec/CMakeLists.txt

@@ -25,13 +25,13 @@ atlas_add_library( tauRecLib
                    src/*.cxx
                    PUBLIC_HEADERS tauRec
                    INCLUDE_DIRS ${ROOT_INCLUDE_DIRS}
-                   LINK_LIBRARIES ${ROOT_LIBRARIES} AthenaBaseComps AthenaKernel GaudiKernel InDetReadoutGeometry TRT_ReadoutGeometry ParticleEvent tauRecToolsLib
+                   LINK_LIBRARIES ${ROOT_LIBRARIES} AthenaBaseComps AthenaKernel GaudiKernel InDetReadoutGeometry TRT_ReadoutGeometry ParticleEvent tauRecToolsLib CaloEvent xAODCaloEvent CaloUtilsLib
                    PRIVATE_LINK_LIBRARIES tauEvent )
 
 atlas_add_component( tauRec
                      src/components/*.cxx
                      INCLUDE_DIRS ${ROOT_INCLUDE_DIRS}
-                     LINK_LIBRARIES ${ROOT_LIBRARIES} AthenaBaseComps AthenaKernel GaudiKernel ParticleEvent tauRecToolsLib tauEvent tauRecLib )
+                     LINK_LIBRARIES ${ROOT_LIBRARIES} AthenaBaseComps AthenaKernel GaudiKernel ParticleEvent tauRecToolsLib CaloEvent xAODCaloEvent CaloUtilsLib tauEvent tauRecLib )
 
 # Install files from the package:
 atlas_install_python_modules( python/*.py )

Reconstruction/tauRec/python/TauRecRunConfigured.py

@@ -31,10 +31,10 @@ class TauRecRunConfigured ( Configured ) :
         from tauRec.tauRecFlags import tauFlags
         self._TauRunnerAlgHandle = TauRunnerAlg ( name=self.name+'Alg', 
                                                   Key_tauInputContainer="tmp_TauJets",
-                                                  Key_Pi0ClusterInputContainer="TauInitialPi0Clusters",
+                                                  Key_Pi0ClusterInputContainer="TauPi0Clusters",
                                                   Key_tauOutputContainer="TauJets",
                                                   Key_neutralPFOOutputContainer="TauNeutralParticleFlowObjects",
-                                                  Key_pi0ClusterOutputContainer="TauPi0Clusters",
+                                                  Key_pi0ClusterOutputContainer="TauSelectedPi0Clusters",
                                                   Key_hadronicPFOOutputContainer="TauHadronicParticleFlowObjects",
                                                   Key_vertexOutputContainer = "TauSecondaryVertices",
                                                   Key_chargedPFOOutputContainer = "TauChargedParticleFlowObjects",

Reconstruction/tauRec/share/TauAODList.py

@@ -29,16 +29,11 @@ TauAODList += [ "xAOD::VertexContainer#TauSecondaryVertices" ]
 TauAODList += [ "xAOD::VertexAuxContainer#TauSecondaryVerticesAux.-vxTrackAtVertex" ]
 
 #------------------------------------------------------------------------------
-# Pi0 cells
+# Pi0 clusters and links to cells
 #------------------------------------------------------------------------------
-TauAODList += [ "CaloCellContainer#TauCommonPi0Cells" ]
-TauAODList += [ "CaloClusterCellLinkContainer#TauInitialPi0Clusters_links" ]
-
-#------------------------------------------------------------------------------
-# Pi0 initial clusters
-#------------------------------------------------------------------------------
-TauAODList += [ "xAOD::CaloClusterContainer#TauInitialPi0Clusters" ]
-TauAODList += [ "xAOD::CaloClusterAuxContainer#TauInitialPi0ClustersAux." ]
+TauAODList += [ "xAOD::CaloClusterContainer#TauPi0Clusters" ]
+TauAODList += [ "xAOD::CaloClusterAuxContainer#TauPi0ClustersAux." ]
+TauAODList += [ "CaloClusterCellLinkContainer#TauPi0Clusters_links" ]
 
 #------------------------------------------------------------------------------
 # Shot clusters
@@ -58,12 +53,6 @@ TauAODList += [ "xAOD::ParticleAuxContainer#TauFinalPi0sAux." ]
 TauAODList += [ "xAOD::PFOContainer#TauShotParticleFlowObjects" ]
 TauAODList += [ "xAOD::PFOAuxContainer#TauShotParticleFlowObjectsAux." ]
 
-#------------------------------------------------------------------------------
-# Cell-based charged ParticleFlowObjects
-#------------------------------------------------------------------------------
-TauAODList += [ "xAOD::PFOContainer#TauChargedParticleFlowObjects" ]
-TauAODList += [ "xAOD::PFOAuxContainer#TauChargedParticleFlowObjectsAux." ]
-
 #------------------------------------------------------------------------------
 # Cell-based neutral ParticleFlowObjects
 #------------------------------------------------------------------------------

Reconstruction/tauRec/share/TauESDList.py

@@ -29,16 +29,11 @@ TauESDList += [ "xAOD::VertexContainer#TauSecondaryVertices" ]
 TauESDList += [ "xAOD::VertexAuxContainer#TauSecondaryVerticesAux.-vxTrackAtVertex" ]
 
 #------------------------------------------------------------------------------
-# Pi0 cells
+# Pi0 clusters and links to cells
 #------------------------------------------------------------------------------
-TauESDList += [ "CaloCellContainer#TauCommonPi0Cells" ]
-TauESDList += [ "CaloClusterCellLinkContainer#TauInitialPi0Clusters_links" ]
-
-#------------------------------------------------------------------------------
-# Pi0 initial clusters
-#------------------------------------------------------------------------------
-TauESDList += [ "xAOD::CaloClusterContainer#TauInitialPi0Clusters" ]
-TauESDList += [ "xAOD::CaloClusterAuxContainer#TauInitialPi0ClustersAux." ]
+TauESDList += [ "xAOD::CaloClusterContainer#TauPi0Clusters" ]
+TauESDList += [ "xAOD::CaloClusterAuxContainer#TauPi0ClustersAux." ]
+TauESDList += [ "CaloClusterCellLinkContainer#TauPi0Clusters_links" ]
 
 #------------------------------------------------------------------------------
 # Shot clusters

Reconstruction/tauRec/share/tauRec_jobOptions.py

@@ -13,14 +13,27 @@ from tauRec.tauRecFlags import tauFlags
 # use Tau Jet Vertex Association Tool
 _doTJVA = tauFlags.doTJVA()
 
-# Pi0 cell RoI making and topoclustering
+# Pi0 RoI making and topoclustering
 _doPi0Clus = tauFlags.doPi0()
 
 from tauRec.TauRecBuilder import TauRecCoreBuilder
 TauRecCoreBuilder(doPi0Clus=_doPi0Clus, doTJVA=_doTJVA)
 
 if _doPi0Clus:
+    # run pi0 topoclustering using cells from TauCommonPi0Cells
     include("tauRec/Pi0ClusterMaker_jobOptions.py")
 
+    # recreate pi0 clusters using cells from AllCalo
+    from tauRec.tauRecConf import ClusterCellRelinkAlg 
+    alg = ClusterCellRelinkAlg('ClusterCellRelinkAlg',
+                               Cells = 'AllCalo',
+                               ClustersInput = 'TauInitialPi0Clusters',
+                               ClustersOutput = 'TauPi0Clusters',
+                               CellLinksOutput = 'TauPi0Clusters_links')
+
+    from AthenaCommon.AlgSequence import AlgSequence
+    topSequence = AlgSequence()
+    topSequence += alg
+
 from tauRec.TauRecRunner import TauRecRunner
 TauRecRunner()

Reconstruction/tauRec/src/ClusterCellRelinkAlg.cxx

+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration.
+*/
+
+#include "tauRec/ClusterCellRelinkAlg.h"
+#include "StoreGate/ReadHandle.h"
+#include "StoreGate/WriteHandle.h"
+#include "xAODCaloEvent/CaloClusterAuxContainer.h"
+#include "CaloUtils/CaloClusterStoreHelper.h"
+
+/**
+ * @brief Gaudi initialize method.
+ */
+StatusCode ClusterCellRelinkAlg::initialize()
+{
+  ATH_CHECK( m_cells.initialize() );
+  ATH_CHECK( m_clustersInput.initialize() );
+  ATH_CHECK( m_clustersOutput.initialize() );
+  ATH_CHECK( m_cellLinksOutput.initialize() );
+
+  return StatusCode::SUCCESS;
+}
+
+/**
+ * @brief Execute the algorithm.
+ * @param ctx Current event context.
+ */
+StatusCode ClusterCellRelinkAlg::execute (const EventContext& ctx) const
+{
+  SG::ReadHandle<CaloCellContainer> cells (m_cells, ctx);
+  if(!cells.isValid()) {
+    ATH_MSG_WARNING( "Collection " << m_cells.key() << " is not valid" );
+    return StatusCode::SUCCESS;
+  }
+  
+  // input clusters
+  SG::ReadHandle<xAOD::CaloClusterContainer> clustersInput (m_clustersInput, ctx);
+  if(!clustersInput.isValid()) {
+    ATH_MSG_WARNING( "Collection " << m_clustersInput.key() << " is not valid" );
+    return StatusCode::SUCCESS;
+  }
+
+  // create new clusters made of cells from another container
+  SG::WriteHandle<xAOD::CaloClusterContainer> clustersOutputHandle (m_clustersOutput, ctx);
+  ATH_CHECK( clustersOutputHandle.record(std::make_unique<xAOD::CaloClusterContainer>(), std::make_unique<xAOD::CaloClusterAuxContainer>()) );
+  xAOD::CaloClusterContainer* clustersOutputContainer = clustersOutputHandle.ptr();
+  clustersOutputContainer->reserve(clustersInput->size());
+
+  for (const xAOD::CaloCluster* oldCluster : *clustersInput) { 
+    xAOD::CaloCluster* newCluster = new xAOD::CaloCluster();
+    clustersOutputContainer->push_back(newCluster);
+    *newCluster = *oldCluster;
+    newCluster->addCellLink(new CaloClusterCellLink(cells.ptr()));
+
+    const CaloClusterCellLink* cellLinks = oldCluster->getCellLinks();
+    CaloClusterCellLink::const_iterator lnk_it=cellLinks->begin();
+    CaloClusterCellLink::const_iterator lnk_it_e=cellLinks->end();
+    newCluster->getOwnCellLinks()->reserve(cellLinks->size());
+
+    for (;lnk_it!=lnk_it_e;++lnk_it) {
+      const IdentifierHash cellHash = (*lnk_it)->caloDDE()->calo_hash();
+      const int newIdx=cells->findIndex(cellHash);
+      newCluster->addCell(newIdx, lnk_it.weight());
+    }
+  }
+
+  SG::WriteHandle<CaloClusterCellLinkContainer> cellLinksOutputHandle (m_cellLinksOutput, ctx);
+  ATH_CHECK( CaloClusterStoreHelper::finalizeClusters(cellLinksOutputHandle, clustersOutputContainer) );
+
+  return StatusCode::SUCCESS;
+}

Reconstruction/tauRec/src/TauCellThinningAlg.cxx

+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration.
+*/
+
+#include "tauRec/TauCellThinningAlg.h"
+#include "StoreGate/ReadHandle.h"
+#include "StoreGate/ThinningHandle.h"
+#include "xAODCaloEvent/CaloCluster.h"
+#include "tauRecTools/HelperFunctions.h"
+#include <vector>
+
+/**
+ * @brief Gaudi initialize method.
+ */
+StatusCode TauCellThinningAlg::initialize()
+{
+  ATH_CHECK( m_cells.initialize(m_streamName) );
+  ATH_CHECK( m_taus.initialize() );
+
+  return StatusCode::SUCCESS;
+}
+
+/**
+ * @brief Execute the algorithm.
+ * @param ctx Current event context.
+ */
+StatusCode TauCellThinningAlg::execute (const EventContext& ctx) const
+{
+  SG::ThinningHandle<CaloCellContainer> cells (m_cells, ctx);
+  cells.thinAll();
+
+  SG::ReadHandle<xAOD::TauJetContainer> taus (m_taus, ctx);
+  if (!taus.isValid()) {
+    ATH_MSG_WARNING( "Collection " << m_taus.key() << " is not valid" );
+    return StatusCode::SUCCESS;
+  }
+
+  for (const xAOD::TauJet* tau : *taus) {
+    std::vector<const xAOD::CaloCluster*> clusters;
+    ATH_CHECK( tauRecTools::GetJetClusterList(tau->jet(), clusters, m_useSubtractedCluster) );
+
+    for (const xAOD::CaloCluster* cluster : clusters) {
+      // only keep cells for clusters within 0.2
+      if (cluster->p4().DeltaR(tau->p4())>0.2) continue;
+
+      const CaloClusterCellLink* cellLinks = cluster->getCellLinks();
+      if (!cellLinks) {
+	ATH_MSG_WARNING( "Cluster without cell links found! Cells cannot be written in AOD." );
+	continue;
+      }
+      
+      if (cellLinks->getCellContainerLink().dataID() != m_cells.key()) {
+	ATH_MSG_WARNING( "Cluster points to cell container "
+			 << cellLinks->getCellContainerLink().dataID()
+			 << " which is different from the cell container being thinned: "
+			 << m_cells.key() << "; cluster skipped.");
+	continue;
+      }
+      
+      CaloClusterCellLink::const_iterator it = cellLinks->begin();
+      CaloClusterCellLink::const_iterator end = cellLinks->end();
+      for (; it != end; ++it) {
+	if (it.index() >= cells->size()) {
+	  ATH_MSG_WARNING( "  Cell index " << it.index()
+			   << " is larger than the number of cells in "
+			   << m_cells.key() << " (" << cells->size() << ")" );
+	  continue;
+	}
+	cells.keep (it.index());
+      }
+    }
+  }
+  
+  return StatusCode::SUCCESS;
+}

Reconstruction/tauRec/src/components/tauRec_entries.cxx

 #include "tauRec/TauProcessorAlg.h"
 #include "tauRec/TauRunnerAlg.h"
+#include "tauRec/TauCellThinningAlg.h"
+#include "tauRec/ClusterCellRelinkAlg.h"
 
 DECLARE_COMPONENT( TauProcessorAlg )
 DECLARE_COMPONENT( TauRunnerAlg )
+DECLARE_COMPONENT( TauCellThinningAlg )
+DECLARE_COMPONENT( ClusterCellRelinkAlg )

Reconstruction/tauRec/tauRec/ClusterCellRelinkAlg.h

+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration.
+*/
+
+#ifndef TAUREC_CLUSTERCELLRELINK_H
+#define TAUREC_CLUSTERCELLRELINK_H
+
+#include "AthenaBaseComps/AthReentrantAlgorithm.h"
+#include "CaloEvent/CaloCellContainer.h"
+#include "CaloEvent/CaloClusterCellLinkContainer.h"
+#include "xAODCaloEvent/CaloClusterContainer.h"
+#include "StoreGate/ReadHandleKey.h"
+#include "StoreGate/WriteHandleKey.h"
+
+/**
+ * @brief Algorithm that rebuilds clusters using cells from a different container.
+ * Only the cell links are updated, other cluster properties are not recomputed.
+ * The CaloClusterCellLinkContainer pointing to the new CaloCellContainer is also created.
+ */
+class ClusterCellRelinkAlg : public AthReentrantAlgorithm
+{
+  using AthReentrantAlgorithm::AthReentrantAlgorithm;
+  
+ public:
+
+  /**
+   * @brief Gaudi initialize method.
+   */
+  virtual StatusCode initialize() override;
+
+  /**
+   * @brief Execute the algorithm.
+   * @param ctx Current event context.
+   */
+  virtual StatusCode execute(const EventContext& ctx) const override;
+
+
+ private:
+
+  // Cell container the output clusters should point to
+  SG::ReadHandleKey<CaloCellContainer> m_cells
+    { this, "Cells", "", "Cell container to point to." };
+  
+  // Input clusters
+  SG::ReadHandleKey<xAOD::CaloClusterContainer> m_clustersInput
+    { this, "ClustersInput", "", "Input cluster container." };
+  
+  // Output clusters
+  SG::WriteHandleKey<xAOD::CaloClusterContainer> m_clustersOutput
+    { this, "ClustersOutput", "", "Output cluster container." };
+
+  // Cell link container for output clusters
+  SG::WriteHandleKey<CaloClusterCellLinkContainer> m_cellLinksOutput
+    { this,"CellLinksOutput", "", "Output CellLinks container" };
+
+};
+
+#endif // not TAUREC_CLUSTERCELLRELINK_H

Reconstruction/tauRec/tauRec/TauCellThinningAlg.h

+/*
+  Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration.
+*/
+
+#ifndef TAUREC_TAUCELLTHINNING_H
+#define TAUREC_TAUCELLTHINNING_H
+
+#include "AthenaBaseComps/AthReentrantAlgorithm.h"
+#include "CaloEvent/CaloCellContainer.h"
+#include "xAODTau/TauJetContainer.h"
+#include "StoreGate/ReadHandleKey.h"
+#include "StoreGate/ThinningHandleKey.h"
+
+/**
+ * @brief Thin calorimeter cells not associated with tau clusters
+ *
+ * Thinning algorithm to keep calorimeter cells from tau seed-jet topoclusters
+ * Based on CaloThinCellsByClusterAlg from CaloRec package
+ */
+class TauCellThinningAlg : public AthReentrantAlgorithm
+{
+  using AthReentrantAlgorithm::AthReentrantAlgorithm;
+
+ public:
+
+  /**
+   * @brief Gaudi initialize method.
+   */
+  virtual StatusCode initialize() override;
+
+  /**
+   * @brief Execute the algorithm.
+   * @param ctx Current event context.
+   */
+  virtual StatusCode execute(const EventContext& ctx) const override;
+
+
+ private:
+
+  // Name of the stream being thinned
+  StringProperty m_streamName
+    { this, "StreamName", "StreamAOD", "Name of the stream being thinned" };
+
+  // Cell container to thin
+  SG::ThinningHandleKey<CaloCellContainer> m_cells
+    { this, "Cells", "AllCalo", "Cell container to thin" };
+
+  // Tau container
+  SG::ReadHandleKey<xAOD::TauJetContainer> m_taus
+    { this, "Taus", "TauJets", "Container of taus for which cells should be saved" };
+
+  // Use shower subtracted clusters, only relevant for PFlow seed jets
+  Gaudi::Property<bool> m_useSubtractedCluster 
+    { this, "UseSubtractedCluster", false, "Use shower subtracted clusters" };
+
+};
+
+#endif // not TAUREC_TAUCELLTHINNING_H

Reconstruction/tauRecTools/src/TauPi0ClusterCreator.cxx

@@ -42,7 +42,7 @@ StatusCode TauPi0ClusterCreator::initialize() {
 //______________________________________________________________________________
 StatusCode TauPi0ClusterCreator::executePi0ClusterCreator(xAOD::TauJet& pTau, xAOD::PFOContainer& neutralPFOContainer,
 							  xAOD::PFOContainer& hadronicClusterPFOContainer,
-							  xAOD::CaloClusterContainer& pi0CaloClusterContainer,
+							  xAOD::CaloClusterContainer& /*pi0CaloClusterContainer*/,
 							  const xAOD::CaloClusterContainer& pPi0ClusterContainer) const
 {
     // Any tau needs to have PFO vectors. Set empty vectors before nTrack cut
@@ -94,11 +94,16 @@ StatusCode TauPi0ClusterCreator::executePi0ClusterCreator(xAOD::TauJet& pTau, xA
         // (not a copy!) since the pointer will otherwise be different than in clusterToShotMap
         std::vector<unsigned> shotsInCluster = getShotsMatchedToCluster( shotVector, clusterToShotMap, cluster);
 
+	/*
         // Make a copy of the cluster to store in output container.
         xAOD::CaloCluster* pPi0Cluster = new xAOD::CaloCluster( *cluster );
 
         // store pi0 calo cluster in the output container
         pi0CaloClusterContainer.push_back(pPi0Cluster);
+	*/
+
+	// temporary, will be cleaned up once pi0CaloClusterContainer is dropped
+	const xAOD::CaloCluster* pPi0Cluster = cluster;
 
         // Calculate input variables for fake supression. 
         // Do this before applying the vertex correction, 
@@ -155,7 +160,9 @@ StatusCode TauPi0ClusterCreator::executePi0ClusterCreator(xAOD::TauJet& pTau, xA
 
         // Set PFO variables
         ElementLink<xAOD::CaloClusterContainer> clusElementLink;
-        clusElementLink.toContainedElement( pi0CaloClusterContainer, pPi0Cluster );
+	// will be cleaned up once pi0CaloClusterContainer is dropped
+        //clusElementLink.toContainedElement( pi0CaloClusterContainer, pPi0Cluster );
+        clusElementLink.toContainedElement( pPi0ClusterContainer, pPi0Cluster );
         neutralPFO->setClusterLink( clusElementLink );
         
         neutralPFO->setP4( (float) pPi0Cluster->pt(), (float) pPi0Cluster->eta(), (float) pPi0Cluster->phi(), (float) pPi0Cluster->m());