Gaudi/NTuple/Writer.h

0001 /***********************************************************************************\
0002 * (c) Copyright 2024 CERN for the benefit of the LHCb and ATLAS collaborations      *
0003 *                                                                                   *
0004 * This software is distributed under the terms of the Apache version 2 licence,     *
0005 * copied verbatim in the file "LICENSE".                                            *
0006 *                                                                                   *
0007 * In applying this licence, CERN does not waive the privileges and immunities       *
0008 * granted to it by virtue of its status as an Intergovernmental Organization        *
0009 * or submit itself to any jurisdiction.                                             *
0010 \***********************************************************************************/
0011 #pragma once
0012
0013 #include <Gaudi/Algorithm.h>
0014 #include <Gaudi/Functional/Consumer.h>
0015 #include <Gaudi/Interfaces/IFileSvc.h>
0016 #include <Gaudi/details/BranchWrapper.h>
0017 #include <GaudiKernel/StdArrayAsProperty.h>
0018 #include <TFile.h>
0019 #include <TTree.h>
0020 #include <fmt/format.h>
0021 #include <gsl/pointers>
0022 #include <gsl/span>
0023 #include <mutex>
0024 #include <numeric>
0025 #include <tuple>
0026 #include <utility>
0027
0028 namespace Gaudi::NTuple {
0029
0030   /**
0031    * @struct WriterMixin
0032    * @brief  Base template for NTuple::WriterMixin.
0033    * Actual specializations of this template provide the functionality.
0034    */
0035   template <typename Signature>
0036   struct WriterMixin {
0037     WriterMixin() = 0; // This function has to be implemented with this speialization:
0038                        // WriterMixin<std::tuple<OUTPUTs...>( const INPUTs&... )>
0039   };
0040
0041   /**
0042    * @struct WriterMixin
0043    * @brief Specialization of WriterMixin for sets of input and output types.
0044    * Provides the functionality for initializing branches in a ROOT TTree,
0045    * transforming input data, filling the tree, and writing it to a file.
0046    *
0047    * @tparam OUTPUTs Variadic template parameters representing output data types.
0048    * @tparam INPUTs Variadic template parameters representing input data types.
0049    */
0050   template <typename... OUTPUTs, typename... INPUTs>
0051   struct WriterMixin<std::tuple<OUTPUTs...>( const INPUTs&... )> {
0052     TTree*                                        m_tree = nullptr;   // Pointer to the TTree being written to
0053     std::array<std::string, sizeof...( OUTPUTs )> m_branchNames;      // Names of the branches to be created
0054     mutable std::vector<details::BranchWrapper>   m_branchWrappers{}; // Container for BranchWrapper objects
0055
0056     /**
0057      * Transform input data to the desired output format. Must be overridden in derived classes.
0058      * @param inputs Variadic inputs of types specified by INPUTs...
0059      * @return A tuple of output data corresponding to OUTPUTs...
0060      */
0061     virtual std::tuple<OUTPUTs...> transform( const INPUTs&... inputs ) const = 0;
0062
0063     // Initialize the TTree and creates branches
0064     void initTree( const std::shared_ptr<TFile>& file, const gsl::span<std::string, sizeof...( OUTPUTs )> branchNames,
0065                    const Gaudi::Algorithm& algRef ) {
0066       file->cd();
0067       m_tree = std::make_unique<TTree>( "WriterTree", "Tree of Writer Algorithm" ).release();
0068       m_branchWrappers.reserve( m_branchWrappers.size() + sizeof...( OUTPUTs ) );
0069       createBranchesForOutputs( branchNames, std::make_index_sequence<sizeof...( OUTPUTs )>{}, algRef );
0070     }
0071
0072     // Create branches in the TTree based on the provided names and output data types
0073     template <std::size_t... Is>
0074     void createBranchesForOutputs( const gsl::span<std::string, sizeof...( OUTPUTs )> branchNames,
0075                                    const std::index_sequence<Is...>, const Gaudi::Algorithm& algRef ) const {
0076       ( ..., m_branchWrappers.emplace_back(
0077                  m_tree, System::typeinfoName( typeid( std::tuple_element_t<Is, std::tuple<OUTPUTs...>> ) ),
0078                  branchNames[Is], "", algRef ) );
0079     }
0080
0081     // Fill the TTree with transformed data from the input
0082     void fillTree( const INPUTs&... inputs ) const {
0083       auto transformedData = transform( inputs... );
0084       std::apply(
0085           [&]( const auto&... elems ) {
0086             size_t index = 0;
0087             ( ..., m_branchWrappers[index++].setDataPtr( const_cast<void*>( static_cast<const void*>( &elems ) ) ) );
0088           },
0089           transformedData );
0090       m_tree->Fill();
0091     }
0092
0093     // Write the TTree to the associated ROOT file
0094     void writeTree( const std::shared_ptr<TFile>& file, const Gaudi::Algorithm& algRef ) {
0095       file->cd();
0096       if ( m_tree->Write() <= 0 ) {
0097         throw GaudiException( "Failed to write TTree to ROOT file.", algRef.name(), StatusCode::FAILURE );
0098       }
0099       m_tree = nullptr;
0100       algRef.info() << "TTree written to TFile." << endmsg;
0101     }
0102
0103     virtual ~WriterMixin() = default;
0104   };
0105
0106   /**
0107    * @struct Writer
0108    * @brief Base template for NTuple::Writer.
0109    * Actual specializations of this template provide the functionality.
0110    */
0111   template <typename Signature, typename Traits_ = Gaudi::Functional::Traits::BaseClass_t<Gaudi::Algorithm>>
0112   struct Writer {
0113     Writer() = 0; // If you wish to not provide any transformation for your data, please use the NTuple::GenericWriter
0114                   // algorithm
0115   };
0116
0117   /**
0118    * @struct Writer
0119    * @brief Specialized struct for NTuple writing with transformation. Inherits from Gaudi::Functional::Consumer
0120    * for processing input data, and WriterMixin for NTuple-specific functionalities.
0121    *
0122    * @tparam OUTPUTs Variadic template parameters representing output data types.
0123    * @tparam INPUTs Variadic template parameters representing input data types.
0124    */
0125   template <typename... OUTPUTs, typename... INPUTs, typename Traits_>
0126   struct Writer<std::tuple<OUTPUTs...>( const INPUTs&... ), Traits_>
0127       : Gaudi::Functional::Consumer<void( const INPUTs&... ), Traits_>,
0128         WriterMixin<std::tuple<OUTPUTs...>( const INPUTs&... )> {
0129     using Consumer_t = Gaudi::Functional::Consumer<void( const INPUTs&... ), Traits_>;
0130     using Consumer_t::Consumer_t;
0131
0132     Gaudi::Property<std::string> m_fileId{ this, "OutputFile", "NTuple", "Identifier for the TFile to write to." };
0133     Gaudi::Property<std::array<std::string, sizeof...( OUTPUTs )>> m_branchNames{
0134         this, "BranchNames", {}, "Names of the tree branches." }; // Names for the tree branches
0135     std::shared_ptr<TFile>       m_file = nullptr;                // Pointer to the ROOT file
0136     Gaudi::Interfaces::IFileSvc* m_fileSvc;
0137     mutable std::mutex           m_mtx; // Mutex for thread-safe operations
0138
0139     // Initialize the algorithm, set up the ROOT file and a TTree branch for each input location
0140     virtual StatusCode initialize() override {
0141       return Consumer_t::initialize().andThen( [this]() {
0142         m_fileSvc = this->template service<Gaudi::Interfaces::IFileSvc>( "FileSvc" );
0143         if ( !m_fileSvc ) {
0144           this->error() << "Failed to retrieve FileSvc." << endmsg;
0145           return StatusCode::FAILURE;
0146         }
0147
0148         m_file = m_fileSvc->getFile( m_fileId );
0149         if ( !m_file ) {
0150           this->error() << "Failed to retrieve TFile." << endmsg;
0151           return StatusCode::FAILURE;
0152         }
0153
0154         this->initTree( m_file, m_branchNames.value(), *this );
0155
0156         return StatusCode::SUCCESS;
0157       } );
0158     }
0159
0160     // Execute the algorithm for each event, retrieving data from the event store and writing it to the TTree
0161     void operator()( const INPUTs&... args ) const override {
0162       std::lock_guard<std::mutex> lock( m_mtx );
0163       this->fillTree( args... );
0164     }
0165
0166     // Finalize the algorithm by writing the TTree to the file and closing it
0167     virtual StatusCode finalize() override {
0168       this->writeTree( m_file, *this );
0169       return Consumer_t::finalize();
0170     }
0171   };
0172 } // namespace Gaudi::NTuple