JANA/Components/JPodioOutput.h

0001 #pragma once
0002 #include <JANA/Utils/JTypeInfo.h>
0003 #include <JANA/JFactorySet.h>
0004 #include <JANA/Components/JHasOutputs.h>
0005 #include <JANA/Components/JPodioDatabundle.h>
0006 #include <JANA/Components/JLightweightDatabundle.h>
0007 #include <podio/Frame.h>
0008 #include <memory>
0009
0010
0011 namespace jana::components {
0012
0013
0014 template <typename PodioT>
0015 class PodioOutput : public JHasOutputs::OutputBase {
0016 private:
0017     std::unique_ptr<typename PodioT::collection_type> m_transient_collection;
0018     bool m_is_subset = false;
0019     JPodioDatabundle* m_podio_databundle;
0020
0021 public:
0022     PodioOutput(JHasOutputs* owner, std::string collection_name="") {
0023
0024         owner->RegisterOutput(this);
0025         this->m_podio_databundle = new JPodioDatabundle;
0026         SetDatabundle(m_podio_databundle);
0027
0028         m_podio_databundle->SetUniqueName(collection_name);
0029         m_podio_databundle->SetTypeName(JTypeInfo::demangle<PodioT>());
0030         m_podio_databundle->SetTypeIndex(std::type_index(typeid(PodioT)));
0031
0032         m_transient_collection = std::move(std::make_unique<typename PodioT::collection_type>());
0033     }
0034
0035     std::unique_ptr<typename PodioT::collection_type>& operator()() { return m_transient_collection; }
0036
0037     typename PodioT::collection_type* operator->() { return m_transient_collection.get(); }
0038
0039     JPodioDatabundle* GetDatabundle() const { return m_podio_databundle; }
0040
0041     void SetSubsetCollection(bool is_subset) {
0042         m_is_subset = is_subset;
0043         m_transient_collection->setSubsetCollection(is_subset);
0044     }
0045
0046     bool IsSubsetCollection() const { return m_is_subset; }
0047
0048
0049 protected:
0050     void ClearData() {
0051         m_podio_databundle->ClearData();
0052     }
0053
0054     void LagrangianStore(JFactorySet& facset, JDatabundle::Status status) override {
0055
0056         auto* bundle = facset.GetDatabundle("podio::Frame");
0057         JLightweightDatabundleT<podio::Frame>* typed_bundle = nullptr;
0058
0059         if (bundle == nullptr) {
0060             // LOG << "No frame databundle found. Creating new databundle.";
0061             typed_bundle = new JLightweightDatabundleT<podio::Frame>;
0062             facset.Add(typed_bundle);
0063         }
0064         else {
0065             typed_bundle = dynamic_cast<JLightweightDatabundleT<podio::Frame>*>(bundle);
0066             if (typed_bundle == nullptr) {
0067                 throw JException("Databundle with unique_name 'podio::Frame' is not a JLightweightDatabundleT");
0068             }
0069         }
0070         if (typed_bundle->GetSize() == 0) {
0071             // LOG << "Found typed bundle with no frame. Creating new frame.";
0072             typed_bundle->GetData().push_back(new podio::Frame);
0073             typed_bundle->SetStatus(JDatabundle::Status::Inserted);
0074         }
0075         podio::Frame* frame = typed_bundle->GetData().at(0);
0076
0077         // LOG << "Storing podio collection with name=" << m_podio_databundle->GetUniqueName() << " to frame " << frame << "...";
0078         frame->put(std::move(m_transient_collection), m_podio_databundle->GetUniqueName());
0079         // LOG << "...done";
0080         const auto* moved = &frame->template get<typename PodioT::collection_type>(m_podio_databundle->GetUniqueName());
0081         m_podio_databundle->SetCollection(moved);
0082         m_podio_databundle->SetStatus(status);
0083         m_transient_collection = std::make_unique<typename PodioT::collection_type>();
0084     }
0085
0086
0087     void EulerianStore(JFactorySet& facset) override {
0088
0089         // First we retrieve the podio::Frame
0090
0091         auto* frame_bundle = facset.GetDatabundle("podio::Frame");
0092         JLightweightDatabundleT<podio::Frame>* typed_frame_bundle = nullptr;
0093
0094         if (frame_bundle == nullptr) {
0095             // LOG << "No frame databundle found. Creating new databundle.";
0096             typed_frame_bundle = new JLightweightDatabundleT<podio::Frame>;
0097             facset.Add(typed_frame_bundle);
0098         }
0099         else {
0100             typed_frame_bundle = dynamic_cast<JLightweightDatabundleT<podio::Frame>*>(frame_bundle);
0101             if (typed_frame_bundle == nullptr) {
0102                 throw JException("Databundle with unique_name 'podio::Frame' is not a JLightweightDatabundleT");
0103             }
0104         }
0105         if (typed_frame_bundle->GetSize() == 0) {
0106             // LOG << "Found typed bundle with no frame. Creating new frame.";
0107             typed_frame_bundle->GetData().push_back(new podio::Frame);
0108             typed_frame_bundle->SetStatus(JDatabundle::Status::Inserted);
0109         }
0110         podio::Frame* frame = typed_frame_bundle->GetData().at(0);
0111         frame->put(std::move(m_transient_collection), m_podio_databundle->GetUniqueName());
0112         const auto* published = &frame->template get<typename PodioT::collection_type>(m_podio_databundle->GetUniqueName());
0113
0114         // Then we store the collection itself
0115
0116         JPodioDatabundle* typed_collection_bundle = nullptr;
0117         auto collection_bundle = facset.GetDatabundle(m_podio_databundle->GetTypeIndex(), m_podio_databundle->GetUniqueName());
0118
0119         if (collection_bundle == nullptr) {
0120             // No databundle present. In this case we simply use m_podio_databundle as a template
0121             typed_collection_bundle = new JPodioDatabundle(*m_podio_databundle);
0122             facset.Add(typed_collection_bundle);
0123         }
0124         else {
0125             typed_collection_bundle = dynamic_cast<JPodioDatabundle*>(collection_bundle);
0126             if (typed_collection_bundle == nullptr) {
0127                 // Wrong databundle present
0128                 throw JException("Databundle with unique_name '%s' is not a JPodioDatabundle", m_podio_databundle->GetUniqueName().c_str());
0129             }
0130         }
0131         typed_collection_bundle->SetCollection(published);
0132         typed_collection_bundle->SetStatus(JDatabundle::Status::Inserted);
0133
0134         auto fac = typed_collection_bundle->GetFactory();
0135         if (fac != nullptr) {
0136             UpdateFactoryStatusOnEulerianStore(fac);
0137         }
0138
0139         m_transient_collection = std::make_unique<typename PodioT::collection_type>();
0140         m_transient_collection->setSubsetCollection(m_is_subset);
0141     }
0142 };
0143
0144
0145 template <typename PodioT>
0146 class VariadicPodioOutput : public JHasOutputs::VariadicOutputBase {
0147 private:
0148     std::vector<std::unique_ptr<typename PodioT::collection_type>> m_transient_collections;
0149     bool m_is_subset = false;
0150
0151 public:
0152     VariadicPodioOutput(JHasOutputs* owner, std::vector<std::string> default_collection_names={}) {
0153         owner->RegisterOutput(this);
0154         for (const std::string& name : default_collection_names) {
0155             auto databundle = new JPodioDatabundle;
0156             databundle->SetUniqueName(name);
0157             databundle->SetTypeName(JTypeInfo::demangle<PodioT>());
0158             databundle->SetTypeIndex(std::type_index(typeid(PodioT)));
0159             GetDatabundles().push_back(databundle);
0160             m_transient_collections.push_back(std::make_unique<typename PodioT::collection_type>());
0161         }
0162     }
0163
0164     void SetShortNames(std::vector<std::string> short_names) override {
0165         for (auto* db : GetDatabundles()) {
0166             delete db;
0167         }
0168         GetDatabundles().clear();
0169         m_transient_collections.clear();
0170         for (const std::string& name : short_names) {
0171             auto databundle = new JPodioDatabundle;
0172             databundle->SetShortName(name);
0173             databundle->SetTypeName(JTypeInfo::demangle<PodioT>());
0174             databundle->SetTypeIndex(std::type_index(typeid(PodioT)));
0175             GetDatabundles().push_back(databundle);
0176             m_transient_collections.push_back(std::make_unique<typename PodioT::collection_type>());
0177             m_transient_collections.back()->setSubsetCollection(m_is_subset);
0178         }
0179     }
0180
0181     void SetUniqueNames(std::vector<std::string> unique_names) override {
0182         for (auto* db : GetDatabundles()) {
0183             delete db;
0184         }
0185         GetDatabundles().clear();
0186         m_transient_collections.clear();
0187         for (const std::string& name : unique_names) {
0188             auto databundle = new JPodioDatabundle;
0189             LOG << "SetUniqueNames: Adding databundle with unique_name " << name;
0190             databundle->SetUniqueName(name);
0191             databundle->SetTypeName(JTypeInfo::demangle<PodioT>());
0192             databundle->SetTypeIndex(std::type_index(typeid(PodioT)));
0193             GetDatabundles().push_back(databundle);
0194             m_transient_collections.push_back(std::make_unique<typename PodioT::collection_type>());
0195             m_transient_collections.back()->setSubsetCollection(m_is_subset);
0196         }
0197     }
0198
0199     std::vector<std::unique_ptr<typename PodioT::collection_type>>& operator()() { return m_transient_collections; }
0200
0201     void SetSubsetCollection(bool is_subset) {
0202         m_is_subset = is_subset;
0203         for (auto& coll : m_transient_collections) {
0204             coll->setSubsetCollection(is_subset);
0205         }
0206     }
0207
0208     bool IsSubsetCollection() const { return m_is_subset; }
0209
0210
0211     void LagrangianStore(JFactorySet& facset, JDatabundle::Status status) override {
0212         if (m_transient_collections.size() != GetDatabundles().size()) {
0213             throw JException("VariadicPodioOutput::LagrangianStore() failed: Declared %d collections, but provided %d.", GetDatabundles().size(), m_transient_collections.size());
0214         }
0215         auto* bundle = facset.GetDatabundle("podio::Frame");
0216         JLightweightDatabundleT<podio::Frame>* typed_bundle = nullptr;
0217
0218         if (bundle == nullptr) {
0219             typed_bundle = new JLightweightDatabundleT<podio::Frame>;
0220             facset.Add(typed_bundle);
0221         }
0222         else {
0223             typed_bundle = dynamic_cast<JLightweightDatabundleT<podio::Frame>*>(bundle);
0224         }
0225         if (typed_bundle->GetSize() == 0) {
0226             typed_bundle->GetData().push_back(new podio::Frame);
0227             typed_bundle->SetStatus(JDatabundle::Status::Inserted);
0228         }
0229         podio::Frame* frame = typed_bundle->GetData().at(0);
0230
0231         size_t i = 0;
0232         for (auto& collection : m_transient_collections) {
0233             frame->put(std::move(collection), GetDatabundles()[i]->GetUniqueName());
0234             const auto* moved = &frame->template get<typename PodioT::collection_type>(GetDatabundles()[i]->GetUniqueName());
0235             const auto &databundle = dynamic_cast<JPodioDatabundle*>(GetDatabundles()[i]);
0236             databundle->SetCollection(moved);
0237             databundle->SetStatus(status);
0238             i += 1;
0239             collection = std::make_unique<typename PodioT::collection_type>();
0240             collection->setSubsetCollection(m_is_subset);
0241         }
0242     }
0243
0244     void EulerianStore(JFactorySet& facset) override {
0245
0246         // First we retrieve the podio::Frame
0247
0248         auto* frame_bundle = facset.GetDatabundle("podio::Frame");
0249         JLightweightDatabundleT<podio::Frame>* typed_frame_bundle = nullptr;
0250
0251         if (frame_bundle == nullptr) {
0252             // LOG << "No frame databundle found. Creating new databundle.";
0253             typed_frame_bundle = new JLightweightDatabundleT<podio::Frame>;
0254             facset.Add(typed_frame_bundle);
0255         }
0256         else {
0257             typed_frame_bundle = dynamic_cast<JLightweightDatabundleT<podio::Frame>*>(frame_bundle);
0258             if (typed_frame_bundle == nullptr) {
0259                 throw JException("Databundle with unique_name 'podio::Frame' is not a JLightweightDatabundleT");
0260             }
0261         }
0262         if (typed_frame_bundle->GetSize() == 0) {
0263             // LOG << "Found typed bundle with no frame. Creating new frame.";
0264             typed_frame_bundle->GetData().push_back(new podio::Frame);
0265             typed_frame_bundle->SetStatus(JDatabundle::Status::Inserted);
0266         }
0267         podio::Frame* frame = typed_frame_bundle->GetData().at(0);
0268
0269
0270         int i=0;
0271         for (auto& collection : m_transient_collections) {
0272             frame->put(std::move(collection), GetDatabundles()[i]->GetUniqueName());
0273             const auto* moved = &frame->template get<typename PodioT::collection_type>(GetDatabundles()[i]->GetUniqueName());
0274
0275             JPodioDatabundle* typed_collection_bundle = nullptr;
0276             auto collection_bundle = facset.GetDatabundle(GetDatabundles()[i]->GetTypeIndex(), GetDatabundles()[i]->GetUniqueName());
0277
0278             if (collection_bundle == nullptr) {
0279                 // No databundle present. In this case we create it, using m_databundles[i] as a template
0280                 auto typed_prototype = static_cast<JPodioDatabundle*>(GetDatabundles()[i]);
0281                 typed_collection_bundle = new JPodioDatabundle(*typed_prototype);
0282                 facset.Add(typed_collection_bundle);
0283             }
0284             else {
0285                 typed_collection_bundle = dynamic_cast<JPodioDatabundle*>(collection_bundle);
0286                 if (typed_collection_bundle == nullptr) {
0287                     // Wrong databundle present
0288                     throw JException("Databundle with unique_name '%s' is not a JPodioDatabundle", GetDatabundles()[i]->GetUniqueName().c_str());
0289                 }
0290             }
0291
0292             // Then we store the collection itself
0293             typed_collection_bundle->SetCollection(moved);
0294             typed_collection_bundle->SetStatus(JDatabundle::Status::Inserted);
0295
0296             auto fac = typed_collection_bundle->GetFactory();
0297             if (fac != nullptr) {
0298                 UpdateFactoryStatusOnEulerianStore(fac);
0299             }
0300
0301             // Replace the transient collection
0302             collection = std::make_unique<typename PodioT::collection_type>();
0303             collection->setSubsetCollection(m_is_subset);
0304             i += 1;
0305         }
0306     }
0307 };
0308
0309 } // namespace jana::components
0310
0311 template <typename T> using PodioOutput = jana::components::PodioOutput<T>;
0312 template <typename T> using VariadicPodioOutput = jana::components::VariadicPodioOutput<T>;
0313