From 80c00ea03f766300806ae77d9d9c75a4e9013a3c Mon Sep 17 00:00:00 2001
From: James Richard Catmore <james.catmore@cern.ch>
Date: Thu, 2 Feb 2023 15:38:34 +0100
Subject: [PATCH] Read/Write/Decor handles for TruthDecayCollectionMaker

This is the first of many MRs that will be needed to enable the derivation framework to run in AthenaMT, since many of the tools (especially in the MC truth domain) don't use R/W/D handles. To avoid a monster MR I'm going to migrate them one tool at a time.

Although this necessarily touches C++ it only impacts the reading/writing so will have no effect on the output of the tool. I have tested this on PHYS in serial but a test with MT will require all tools to be migrated.
---
 .../TruthDecayCollectionMaker.h               | 36 ++++++--
 .../src/TruthDecayCollectionMaker.cxx         | 86 ++++++++++++-------
 2 files changed, 83 insertions(+), 39 deletions(-)

diff --git a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/DerivationFrameworkMCTruth/TruthDecayCollectionMaker.h b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/DerivationFrameworkMCTruth/TruthDecayCollectionMaker.h
index 3a9c546f207..cb8d2a55a0b 100644
--- a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/DerivationFrameworkMCTruth/TruthDecayCollectionMaker.h
+++ b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/DerivationFrameworkMCTruth/TruthDecayCollectionMaker.h
@@ -1,5 +1,5 @@
 /*
-  Copyright (C) 2002-2017 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
 */
 
 #ifndef DERIVATIONFRAMEWORK_TRUTHDECAYCOLLECTIONMAKER_H
@@ -11,11 +11,17 @@
 // EDM -- typedefs so these are includes
 #include "xAODTruth/TruthParticleContainer.h"
 #include "xAODTruth/TruthVertexContainer.h"
+// Handles and keys
+#include "StoreGate/WriteDecorHandleKey.h"
+#include "StoreGate/ReadHandleKey.h"
+#include "StoreGate/WriteHandleKey.h"
+#include "StoreGate/WriteHandle.h"
 // Standard library includes
 #include <vector>
 #include <string>
 
 namespace DerivationFramework {
+ 
 
   class TruthDecayCollectionMaker : public AthAlgTool, public IAugmentationTool {
     public: 
@@ -30,15 +36,33 @@ namespace DerivationFramework {
       bool m_keepCHadrons; //!< Option to keep all c-hadrons (better than giving PDG IDs)
       bool m_keepBSM; //!< Option to keep all BSM particles (better than giving PDG IDs)
       bool m_rejectHadronChildren; //!< Option to reject hadron descendants
-      std::string m_particlesKey; //!< Input particle collection (navigates to the vertices)
+      SG::ReadHandleKey<xAOD::TruthParticleContainer> m_particlesKey
+         {this, "ParticlesKey", "TruthParticles", "ReadHandleKey for input TruthParticleContainer"};
+      SG::WriteHandleKey<xAOD::TruthParticleContainer> m_outputParticlesKey
+         {this, "NewParticleKey", "", "WriteHandleKey for new TruthParticleContainer"};
+      SG::WriteHandleKey<xAOD::TruthVertexContainer> m_outputVerticesKey
+         {this, "NewVertexKey", "", "WriteHandleKey for new TruthVertexContainer"};
+      SG::WriteDecorHandleKey<xAOD::TruthParticleContainer> m_originDecoratorKey
+         {this, "classifierParticleOrigin", "TruthParticles.classifierParticleOrigin","Name of the decoration which records the particle origin as determined by the MCTruthClassifier"};
+      SG::WriteDecorHandleKey<xAOD::TruthParticleContainer> m_typeDecoratorKey
+         {this, "classifierParticleType", "TruthParticles.classifierParticleType","Name of the decoration which records the particle type as determined by the MCTruthClassifier"};
+      SG::WriteDecorHandleKey<xAOD::TruthParticleContainer> m_outcomeDecoratorKey
+         {this, "classifierParticleOutCome", "TruthParticles.classifierParticleOutCome","Name of the decoration which records the particle outcome as determined by the MCTruthClassifier"};      
+      SG::WriteDecorHandleKey<xAOD::TruthParticleContainer> m_classificationDecoratorKey
+         {this, "Classification", "TruthParticles.Classification","Name of the decoration which records the particle outcome as determined by the MCTruthClassifier"};
+      SG::WriteDecorHandleKey<xAOD::TruthParticleContainer> m_motherIDDecoratorKey
+         {this, "motherID", "TruthParticles.motherID","Name of the decoration which records the ID of the particle's mother"};
+      SG::WriteDecorHandleKey<xAOD::TruthParticleContainer> m_daughterIDDecoratorKey
+         {this, "daughterID", "TruthParticles.daughterID","Name of the decoration which records the ID of the particle's daughter"};
+
       std::string m_collectionName; //!< Output collection name stem
       int m_generations; //!< Number of generations after the particle in question to keep
       // Helper functions for building up the decay product collections
-      int addTruthParticle( const xAOD::TruthParticle& old_part, xAOD::TruthParticleContainer* part_cont,
-                                  xAOD::TruthVertexContainer* vert_cont, std::vector<int>& seen_particles,
+      int addTruthParticle( const xAOD::TruthParticle& old_part, SG::WriteHandle<xAOD::TruthParticleContainer>& part_cont,
+                            SG::WriteHandle<xAOD::TruthVertexContainer>& vert_cont, std::vector<int>& seen_particles,
                             const int generations=-1) const;
-      int addTruthVertex( const xAOD::TruthVertex& old_vert, xAOD::TruthParticleContainer* part_cont,
-                                xAOD::TruthVertexContainer* vert_cont, std::vector<int>& seen_particles,
+      int addTruthVertex( const xAOD::TruthVertex& old_vert, SG::WriteHandle<xAOD::TruthParticleContainer>& part_cont,
+                          SG::WriteHandle<xAOD::TruthVertexContainer>& vert_cont, std::vector<int>& seen_particles,
                           const int generations=-1) const;
       bool id_ok( const xAOD::TruthParticle& part ) const;
   }; 
diff --git a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthDecayCollectionMaker.cxx b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthDecayCollectionMaker.cxx
index 7589802ea16..2ac95bd36eb 100644
--- a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthDecayCollectionMaker.cxx
+++ b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthDecayCollectionMaker.cxx
@@ -1,5 +1,5 @@
 /*
-  Copyright (C) 2002-2021 CERN for the benefit of the ATLAS collaboration
+  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
 */
 
 /////////////////////////////////////////////////////////////////
@@ -8,6 +8,9 @@
 // Based on TruthCollectionMaker (but simpler)
 
 #include "DerivationFrameworkMCTruth/TruthDecayCollectionMaker.h"
+#include "StoreGate/ReadHandle.h"
+#include "StoreGate/WriteHandle.h"
+#include "StoreGate/WriteDecorHandle.h"
 #include "xAODTruth/TruthParticleContainer.h"
 #include "xAODTruth/TruthParticleAuxContainer.h"
 #include "xAODTruth/TruthVertexContainer.h"
@@ -46,21 +49,35 @@ StatusCode DerivationFramework::TruthDecayCollectionMaker::initialize()
 {
     ATH_MSG_VERBOSE("initialize() ...");
 
-    if (m_particlesKey.empty()) {
-        ATH_MSG_FATAL("No truth particle collection provided to use as a basis for new collections");
-        return StatusCode::FAILURE;
-    } else {ATH_MSG_INFO("Using " << m_particlesKey << " as the source collections for new truth collections");}
+    // Input truth particles
+    ATH_CHECK( m_particlesKey.initialize() );
+    ATH_MSG_INFO("Using " << m_particlesKey.key() << " as the input truth container key");
 
+    // Output particle/vertex containers
     if (m_collectionName.empty()) {
-        ATH_MSG_FATAL("No key provided for the new truth particle collections");
+        ATH_MSG_FATAL("No base name provided for the new truth particle/vertex containers");
         return StatusCode::FAILURE;
-    } else {ATH_MSG_INFO("New truth particle collection key: " << m_collectionName );}
+    } else {ATH_MSG_INFO("Base name for new truth particle/vertex containers: " << m_collectionName );}
+    m_outputParticlesKey = std::string(m_collectionName) + "Particles";
+    ATH_CHECK(m_outputParticlesKey.initialize());
+    ATH_MSG_INFO("New truth particles container key: " << m_outputParticlesKey.key() );
+    m_outputVerticesKey = std::string(m_collectionName) + "Vertices"; 
+    ATH_CHECK(m_outputVerticesKey.initialize());
+    ATH_MSG_INFO("New truth vertices container key: " << m_outputVerticesKey.key() );
 
     if (m_pdgIdsToKeep.empty() && !m_keepBHadrons && !m_keepCHadrons && !m_keepBSM) {
         ATH_MSG_FATAL("No PDG IDs provided, not keeping b- or c-hadrons or BSM particles -- what do you want?");
         return StatusCode::FAILURE;
     }
 
+    // Decorators
+    ATH_CHECK(m_originDecoratorKey.initialize());
+    ATH_CHECK(m_typeDecoratorKey.initialize());
+    ATH_CHECK(m_outcomeDecoratorKey.initialize());
+    ATH_CHECK(m_classificationDecoratorKey.initialize());
+    ATH_CHECK(m_motherIDDecoratorKey.initialize());
+    ATH_CHECK(m_daughterIDDecoratorKey.initialize());
+
     return StatusCode::SUCCESS;
 }
 
@@ -68,34 +85,33 @@ StatusCode DerivationFramework::TruthDecayCollectionMaker::initialize()
 // Selection and collection creation
 StatusCode DerivationFramework::TruthDecayCollectionMaker::addBranches() const
 {
+    // Event context for AthenaMT
+    const EventContext& ctx = Gaudi::Hive::currentContext();
+     
     // Retrieve truth collections
-    const xAOD::TruthParticleContainer* importedTruthParticles(nullptr);
-    if (evtStore()->retrieve(importedTruthParticles,m_particlesKey).isFailure()) {
-        ATH_MSG_ERROR("No TruthParticle collection with name " << m_particlesKey << " found in StoreGate!");
+    SG::ReadHandle<xAOD::TruthParticleContainer> truthParticles(m_particlesKey,ctx);
+    if (!truthParticles.isValid()) {
+        ATH_MSG_ERROR("Couldn't retrieve TruthParticle collection with name " << m_particlesKey);
         return StatusCode::FAILURE;
     }
 
     // Create the new particle containers
-    xAOD::TruthParticleContainer* newParticleCollection = new xAOD::TruthParticleContainer();
-    CHECK( evtStore()->record( newParticleCollection, m_collectionName + "Particles" ) );
-    xAOD::TruthParticleAuxContainer* newParticleAuxCollection = new xAOD::TruthParticleAuxContainer();
-    CHECK( evtStore()->record( newParticleAuxCollection, m_collectionName + "ParticlesAux." ) );
-    newParticleCollection->setStore( newParticleAuxCollection );
-    ATH_MSG_DEBUG( "Recorded new TruthParticleContainer with key: " << (m_collectionName+"Particles"));
+    SG::WriteHandle<xAOD::TruthParticleContainer> newParticleCollection(m_outputParticlesKey, ctx);
+    ATH_CHECK(newParticleCollection.record(std::make_unique<xAOD::TruthParticleContainer>(),
+                                           std::make_unique<xAOD::TruthParticleAuxContainer>()));
+    ATH_MSG_DEBUG( "Recorded new TruthParticleContainer with key: " << (m_outputParticlesKey.key()));
     // Create the new vertex containers
-    xAOD::TruthVertexContainer* newVertexCollection = new xAOD::TruthVertexContainer();
-    CHECK( evtStore()->record( newVertexCollection, m_collectionName + "Vertices" ) );
-    xAOD::TruthVertexAuxContainer* newVertexAuxCollection = new xAOD::TruthVertexAuxContainer();
-    CHECK( evtStore()->record( newVertexAuxCollection, m_collectionName + "VerticesAux." ) );
-    newVertexCollection->setStore( newVertexAuxCollection );
-    ATH_MSG_DEBUG( "Recorded new TruthVertexContainer with key: " << (m_collectionName+"Vertices"));
+    SG::WriteHandle<xAOD::TruthVertexContainer> newVertexCollection(m_outputVerticesKey, ctx);
+    ATH_CHECK(newVertexCollection.record(std::make_unique<xAOD::TruthVertexContainer>(),
+                                         std::make_unique<xAOD::TruthVertexAuxContainer>()));
+    ATH_MSG_DEBUG( "Recorded new TruthVertexContainer with key: " << (m_outputVerticesKey.key()));
 
     // List of barcodes for particles in our collection already.  Because of the way we recurse,
     // adding more particles as we go, there should be no need to add (or help from adding) the
     // barcodes of particles that we are *not* going to keep
     std::vector<int> seen_particles;
     // Go through that list of particles!
-    for (const auto * part : *importedTruthParticles){
+    for (const auto * part : *truthParticles){
         // If this passes my cuts, keep it
         if (id_ok(*part)){
             addTruthParticle( *part, newParticleCollection, newVertexCollection, seen_particles , m_generations );
@@ -104,8 +120,10 @@ StatusCode DerivationFramework::TruthDecayCollectionMaker::addBranches() const
     return StatusCode::SUCCESS;
 }
 
-int DerivationFramework::TruthDecayCollectionMaker::addTruthParticle( const xAOD::TruthParticle& old_part, xAOD::TruthParticleContainer* part_cont, 
-                                                                      xAOD::TruthVertexContainer* vert_cont, std::vector<int>& seen_particles,
+int DerivationFramework::TruthDecayCollectionMaker::addTruthParticle( const xAOD::TruthParticle& old_part, 
+                                                                      SG::WriteHandle<xAOD::TruthParticleContainer>& part_cont, 
+                                                                      SG::WriteHandle<xAOD::TruthVertexContainer>& vert_cont, 
+                                                                      std::vector<int>& seen_particles,
                                                                       const int generations) const {
     // See if we've seen it - note, could also do this with a unary function on the container itself
     if (std::find(seen_particles.begin(),seen_particles.end(),old_part.barcode())!=seen_particles.end()){
@@ -117,12 +135,12 @@ int DerivationFramework::TruthDecayCollectionMaker::addTruthParticle( const xAOD
     // Now we have seen it
     seen_particles.push_back(old_part.barcode());
     // Set up decorators
-    const static SG::AuxElement::Decorator< unsigned int > originDecorator("classifierParticleOrigin");
-    const static SG::AuxElement::Decorator< unsigned int > typeDecorator("classifierParticleType");
-    const static SG::AuxElement::Decorator< unsigned int > outcomeDecorator("classifierParticleOutCome");
-    const static SG::AuxElement::Decorator< unsigned int > classificationDecorator("Classification");
-    const static SG::AuxElement::Decorator< int > motherIDDecorator("motherID");
-    const static SG::AuxElement::Decorator< int > daughterIDDecorator("daughterID");
+    SG::WriteDecorHandle<xAOD::TruthParticleContainer, unsigned int > originDecorator(m_originDecoratorKey);  
+    SG::WriteDecorHandle<xAOD::TruthParticleContainer, unsigned int > typeDecorator(m_typeDecoratorKey);
+    SG::WriteDecorHandle<xAOD::TruthParticleContainer, unsigned int > outcomeDecorator(m_outcomeDecoratorKey);
+    SG::WriteDecorHandle<xAOD::TruthParticleContainer, unsigned int > classificationDecorator(m_classificationDecoratorKey);
+    SG::WriteDecorHandle< xAOD::TruthParticleContainer, int > motherIDDecorator(m_motherIDDecoratorKey);
+    SG::WriteDecorHandle< xAOD::TruthParticleContainer, int > daughterIDDecorator(m_daughterIDDecoratorKey);
     // Make a truth particle and add it to the container
     xAOD::TruthParticle* xTruthParticle = new xAOD::TruthParticle();
     part_cont->push_back( xTruthParticle );
@@ -167,8 +185,10 @@ int DerivationFramework::TruthDecayCollectionMaker::addTruthParticle( const xAOD
     return my_index;
 }
 
-int DerivationFramework::TruthDecayCollectionMaker::addTruthVertex( const xAOD::TruthVertex& old_vert, xAOD::TruthParticleContainer* part_cont, 
-                                                                    xAOD::TruthVertexContainer* vert_cont, std::vector<int>& seen_particles,
+int DerivationFramework::TruthDecayCollectionMaker::addTruthVertex( const xAOD::TruthVertex& old_vert, 
+                                                                    SG::WriteHandle<xAOD::TruthParticleContainer>& part_cont, 
+                                                                    SG::WriteHandle<xAOD::TruthVertexContainer>& vert_cont, 
+                                                                    std::vector<int>& seen_particles,
                                                                     const int generations) const {
     // Make a new vertex and add it to the container
     xAOD::TruthVertex* xTruthVertex = new xAOD::TruthVertex();
-- 
GitLab