File indexing completed on 2025-01-18 09:55:32
0001
0002
0003 #ifndef EDM4EIC_TensorCollection_H
0004 #define EDM4EIC_TensorCollection_H
0005
0006
0007 #include "edm4eic/Tensor.h"
0008 #include "edm4eic/MutableTensor.h"
0009 #include "edm4eic/TensorObj.h"
0010 #include "edm4eic/TensorCollectionData.h"
0011
0012
0013 #include "podio/ICollectionProvider.h"
0014 #include "podio/CollectionBase.h"
0015
0016 #if defined(PODIO_JSON_OUTPUT) && !defined(__CLING__)
0017 #include "nlohmann/json_fwd.hpp"
0018 #endif
0019
0020 #include <string_view>
0021 #include <vector>
0022 #include <array>
0023 #include <algorithm>
0024 #include <ostream>
0025 #include <mutex>
0026 #include <memory>
0027 #include <cstddef>
0028
0029 namespace podio {
0030 struct RelationNames;
0031 }
0032
0033 namespace edm4eic {
0034
0035
0036
0037 class TensorCollectionIterator {
0038 public:
0039 TensorCollectionIterator(size_t index, const TensorObjPointerContainer* collection) : m_index(index), m_object(podio::utils::MaybeSharedPtr<TensorObj>{nullptr}), m_collection(collection) {}
0040
0041 TensorCollectionIterator(const TensorCollectionIterator&) = delete;
0042 TensorCollectionIterator& operator=(const TensorCollectionIterator&) = delete;
0043
0044 bool operator!=(const TensorCollectionIterator& x) const {
0045 return m_index != x.m_index;
0046 }
0047
0048 bool operator==(const TensorCollectionIterator& x) const {
0049 return m_index == x.m_index;
0050 }
0051
0052 Tensor operator*();
0053 Tensor* operator->();
0054 TensorCollectionIterator& operator++();
0055
0056 private:
0057 size_t m_index;
0058 Tensor m_object;
0059 const TensorObjPointerContainer* m_collection;
0060 };
0061
0062
0063 class TensorMutableCollectionIterator {
0064 public:
0065 TensorMutableCollectionIterator(size_t index, const TensorObjPointerContainer* collection) : m_index(index), m_object(podio::utils::MaybeSharedPtr<TensorObj>{nullptr}), m_collection(collection) {}
0066
0067 TensorMutableCollectionIterator(const TensorMutableCollectionIterator&) = delete;
0068 TensorMutableCollectionIterator& operator=(const TensorMutableCollectionIterator&) = delete;
0069
0070 bool operator!=(const TensorMutableCollectionIterator& x) const {
0071 return m_index != x.m_index;
0072 }
0073
0074 bool operator==(const TensorMutableCollectionIterator& x) const {
0075 return m_index == x.m_index;
0076 }
0077
0078 MutableTensor operator*();
0079 MutableTensor* operator->();
0080 TensorMutableCollectionIterator& operator++();
0081
0082 private:
0083 size_t m_index;
0084 MutableTensor m_object;
0085 const TensorObjPointerContainer* m_collection;
0086 };
0087
0088
0089
0090
0091
0092 class TensorCollection : public podio::CollectionBase {
0093 public:
0094 using value_type = Tensor;
0095 using const_iterator = TensorCollectionIterator;
0096 using iterator = TensorMutableCollectionIterator;
0097 using difference_type = ptrdiff_t;
0098 using size_type = size_t;
0099
0100 TensorCollection();
0101 TensorCollection(TensorCollectionData&& data, bool isSubsetColl);
0102
0103 TensorCollection(const TensorCollection& ) = delete;
0104 TensorCollection& operator=(const TensorCollection& ) = delete;
0105 TensorCollection(TensorCollection&&) = default;
0106 TensorCollection& operator=(TensorCollection&&) = default;
0107
0108
0109 ~TensorCollection();
0110
0111 constexpr static auto typeName = "edm4eic::TensorCollection";
0112 constexpr static auto valueTypeName = "edm4eic::Tensor";
0113 constexpr static auto dataTypeName = "edm4eic::TensorData";
0114
0115 void clear() final;
0116
0117
0118 void print(std::ostream& os=std::cout, bool flush=true) const final;
0119
0120
0121 TensorCollection* operator->() { return static_cast<TensorCollection*>(this); }
0122
0123
0124 MutableTensor create();
0125
0126
0127
0128 template<typename... Args>
0129 MutableTensor create(Args&&... args);
0130
0131
0132 std::size_t size() const final;
0133
0134
0135 std::size_t max_size() const final;
0136
0137
0138 bool empty() const final;
0139
0140
0141 const std::string_view getTypeName() const final { return typeName; }
0142
0143 const std::string_view getValueTypeName() const final { return valueTypeName; }
0144
0145 const std::string_view getDataTypeName() const final { return dataTypeName; }
0146
0147 podio::SchemaVersionT getSchemaVersion() const final;
0148
0149 bool isSubsetCollection() const final {
0150 return m_isSubsetColl;
0151 }
0152
0153 void setSubsetCollection(bool setSubset=true) final;
0154
0155
0156 Tensor operator[](std::size_t index) const;
0157
0158 MutableTensor operator[](std::size_t index);
0159
0160 Tensor at(std::size_t index) const;
0161
0162 MutableTensor at(std::size_t index);
0163
0164
0165
0166 void push_back(const MutableTensor& object);
0167
0168 void push_back(const Tensor& object);
0169
0170 void prepareForWrite() const final;
0171 void prepareAfterRead() final;
0172 bool setReferences(const podio::ICollectionProvider* collectionProvider) final;
0173
0174
0175 podio::CollectionWriteBuffers getBuffers() final;
0176
0177 void setID(uint32_t ID) final {
0178 m_collectionID = ID;
0179 if (!m_isSubsetColl) {
0180 std::for_each(m_storage.entries.begin(), m_storage.entries.end(),
0181 [ID] (TensorObj* obj) { obj->id = {obj->id.index, static_cast<uint32_t>(ID)}; }
0182 );
0183 }
0184 m_isValid = true;
0185 }
0186
0187 uint32_t getID() const final {
0188 return m_collectionID;
0189 }
0190
0191 bool isValid() const final {
0192 return m_isValid;
0193 }
0194
0195 size_t getDatamodelRegistryIndex() const final;
0196
0197
0198 iterator begin() {
0199 return iterator(0, &m_storage.entries);
0200 }
0201 const_iterator begin() const {
0202 return const_iterator(0, &m_storage.entries);
0203 }
0204 const_iterator cbegin() const {
0205 return begin();
0206 }
0207 iterator end() {
0208 return iterator(m_storage.entries.size(), &m_storage.entries);
0209 }
0210 const_iterator end() const {
0211 return const_iterator(m_storage.entries.size(), &m_storage.entries);
0212 }
0213 const_iterator cend() const {
0214 return end();
0215 }
0216
0217 std::vector<std::int32_t> elementType(const size_t nElem = 0) const;
0218
0219 private:
0220
0221
0222
0223 friend class TensorCollectionData;
0224
0225 bool m_isValid{false};
0226 mutable bool m_isPrepared{false};
0227 bool m_isSubsetColl{false};
0228 uint32_t m_collectionID{0};
0229 mutable std::unique_ptr<std::mutex> m_storageMtx{nullptr};
0230 mutable TensorCollectionData m_storage{};
0231 };
0232
0233 std::ostream& operator<<(std::ostream& o, const TensorCollection& v);
0234
0235 template<typename... Args>
0236 MutableTensor TensorCollection::create(Args&&... args) {
0237 if (m_isSubsetColl) {
0238 throw std::logic_error("Cannot create new elements on a subset collection");
0239 }
0240 const int size = m_storage.entries.size();
0241 auto obj = new TensorObj({size, m_collectionID}, {std::forward<Args>(args)...});
0242 m_storage.entries.push_back(obj);
0243
0244
0245 obj->m_shape = new std::vector<std::int64_t>();
0246 obj->m_floatData = new std::vector<float>();
0247 obj->m_int64Data = new std::vector<std::int64_t>();
0248 m_storage.createRelations(obj);
0249 return MutableTensor(podio::utils::MaybeSharedPtr(obj));
0250 }
0251
0252 #if defined(PODIO_JSON_OUTPUT) && !defined(__CLING__)
0253 void to_json(nlohmann::json& j, const TensorCollection& collection);
0254 #endif
0255
0256 }
0257
0258
0259 #endif