diff --git a/src/HDF5Reader.cc b/src/HDF5Reader.cc index f310517..de6cd58 100644 --- a/src/HDF5Reader.cc +++ b/src/HDF5Reader.cc @@ -1,305 +1,311 @@ /** * \authors The HEJ collaboration (see AUTHORS for details) * \date 2019 * \copyright GPLv2 or later */ #include "HEJ/HDF5Reader.hh" #ifdef HEJ_BUILD_WITH_HDF5 #include <numeric> #include <iterator> #include "highfive/H5File.hpp" namespace HEJ { namespace { // buffer size for reader // each "reading from disk" reads "chunk_size" many event at once constexpr std::size_t chunk_size = 10000; struct ParticleData { std::vector<int> id; std::vector<int> status; std::vector<int> mother1; std::vector<int> mother2; std::vector<int> color1; std::vector<int> color2; std::vector<double> px; std::vector<double> py; std::vector<double> pz; std::vector<double> e; std::vector<double> m; std::vector<double> lifetime; std::vector<double> spin; }; struct EventRecords { std::vector<int> particle_start; std::vector<int> nparticles; std::vector<int> pid; std::vector<double> weight; + int trials; std::vector<double> scale; std::vector<double> fscale; std::vector<double> rscale; std::vector<double> aqed; std::vector<double> aqcd; ParticleData particles; }; class ConstEventIterator { public: // iterator traits using iterator_category = std::bidirectional_iterator_tag; using value_type = LHEF::HEPEUP; using difference_type = std::ptrdiff_t; using pointer = const LHEF::HEPEUP*; using reference = LHEF::HEPEUP const &; using iterator = ConstEventIterator; friend iterator cbegin(EventRecords const & records) noexcept; friend iterator cend(EventRecords const & records) noexcept; iterator& operator++() { particle_offset_ += records_.get().nparticles[idx_]; ++idx_; return *this; } iterator& operator--() { --idx_; particle_offset_ -= records_.get().nparticles[idx_]; return *this; } iterator operator--(int) { auto res = *this; --(*this); return res; } bool operator==(iterator const & other) const { return idx_ == other.idx_; } bool operator!=(iterator other) const { return !(*this == other); } value_type operator*() const { value_type hepeup{}; auto const & r = records_.get(); hepeup.NUP = r.nparticles[idx_]; hepeup.IDPRUP = r.pid[idx_]; - hepeup.XWGTUP = r.weight[idx_]; + hepeup.XWGTUP = r.weight[idx_]/r.trials; hepeup.weights.emplace_back(hepeup.XWGTUP, nullptr); hepeup.SCALUP = r.scale[idx_]; hepeup.scales.muf = r.fscale[idx_]; hepeup.scales.mur = r.rscale[idx_]; hepeup.AQEDUP = r.aqed[idx_]; hepeup.AQCDUP = r.aqcd[idx_]; const size_t start = particle_offset_; const size_t end = start + hepeup.NUP; auto const & p = r.particles; hepeup.IDUP = std::vector<long>( begin(p.id)+start, begin(p.id)+end ); hepeup.ISTUP = std::vector<int>( begin(p.status)+start, begin(p.status)+end ); hepeup.VTIMUP = std::vector<double>( begin(p.lifetime)+start, begin(p.lifetime)+end ); hepeup.SPINUP = std::vector<double>( begin(p.spin)+start, begin(p.spin)+end ); hepeup.MOTHUP.resize(hepeup.NUP); hepeup.ICOLUP.resize(hepeup.NUP); hepeup.PUP.resize(hepeup.NUP); for(size_t i = 0; i < hepeup.MOTHUP.size(); ++i) { const size_t idx = start + i; assert(idx < end); hepeup.MOTHUP[i] = std::make_pair(p.mother1[idx], p.mother2[idx]); hepeup.ICOLUP[i] = std::make_pair(p.color1[idx], p.color2[idx]); hepeup.PUP[i] = std::vector<double>{ p.px[idx], p.py[idx], p.pz[idx], p.e[idx], p.m[idx] }; } return hepeup; } private: explicit ConstEventIterator(EventRecords const & records): records_{records} {} std::reference_wrapper<const EventRecords> records_; size_t idx_ = 0; size_t particle_offset_ = 0; }; // end ConstEventIterator ConstEventIterator cbegin(EventRecords const & records) noexcept { return ConstEventIterator{records}; } ConstEventIterator cend(EventRecords const & records) noexcept { auto it =ConstEventIterator{records}; it.idx_ = records.aqcd.size(); // or size of any other records member return it; } } // end anonymous namespace struct HDF5Reader::HDF5ReaderImpl{ HighFive::File file; std::size_t event_idx; std::size_t particle_idx; std::size_t nevents; EventRecords records; ConstEventIterator cur_event; LHEF::HEPRUP heprup; LHEF::HEPEUP hepeup; explicit HDF5ReaderImpl(std::string const & filename): file{filename}, event_idx{0}, particle_idx{0}, nevents{ file.getGroup("event") .getDataSet("nparticles") // or any other dataset .getSpace().getDimensions().front() }, records{}, cur_event{cbegin(records)}, heprup{}, hepeup{} { read_heprup(); read_event_records(chunk_size); } void read_heprup() { const auto init = file.getGroup("init"); init.getDataSet( "PDFgroupA" ).read(heprup.PDFGUP.first); init.getDataSet( "PDFgroupB" ).read(heprup.PDFGUP.second); init.getDataSet( "PDFsetA" ).read(heprup.PDFSUP.first); init.getDataSet( "PDFsetB" ).read(heprup.PDFSUP.second); init.getDataSet( "beamA" ).read(heprup.IDBMUP.first); init.getDataSet( "beamB" ).read(heprup.IDBMUP.second); init.getDataSet( "energyA" ).read(heprup.EBMUP.first); init.getDataSet( "energyB" ).read(heprup.EBMUP.second); init.getDataSet( "numProcesses" ).read(heprup.NPRUP); init.getDataSet( "weightingStrategy" ).read(heprup.IDWTUP); const auto proc_info = file.getGroup("procInfo"); proc_info.getDataSet( "procId" ).read(heprup.LPRUP); proc_info.getDataSet( "xSection" ).read(heprup.XSECUP); proc_info.getDataSet( "error" ).read(heprup.XERRUP); // TODO: is this identification correct? proc_info.getDataSet( "unitWeight" ).read(heprup.XMAXUP); + records.trials = 0; + std::vector<int> trials; + file.getGroup("event").getDataSet("trials").read(trials); + for( auto trial: trials) + records.trials+=trial; } std::size_t read_event_records(std::size_t count) { count = std::min(count, nevents-event_idx); auto events = file.getGroup("event"); events.getDataSet("nparticles").select({event_idx}, {count}).read(records.nparticles); assert(records.nparticles.size() == count); events.getDataSet("pid").select( {event_idx}, {count} ).read( records.pid ); events.getDataSet("weight").select( {event_idx}, {count} ).read( records.weight ); events.getDataSet("scale").select( {event_idx}, {count} ).read( records.scale ); events.getDataSet("fscale").select( {event_idx}, {count} ).read( records.fscale ); events.getDataSet("rscale").select( {event_idx}, {count} ).read( records.rscale ); events.getDataSet("aqed").select( {event_idx}, {count} ).read( records.aqed ); events.getDataSet("aqcd").select( {event_idx}, {count} ).read( records.aqcd ); const std::size_t particle_count = std::accumulate( begin(records.nparticles), end(records.nparticles), 0 ); auto pdata = file.getGroup("particle"); auto & particles = records.particles; pdata.getDataSet("id").select( {particle_idx}, {particle_count} ).read( particles.id ); pdata.getDataSet("status").select( {particle_idx}, {particle_count} ).read( particles.status ); pdata.getDataSet("mother1").select( {particle_idx}, {particle_count} ).read( particles.mother1 ); pdata.getDataSet("mother2").select( {particle_idx}, {particle_count} ).read( particles.mother2 ); pdata.getDataSet("color1").select( {particle_idx}, {particle_count} ).read( particles.color1 ); pdata.getDataSet("color2").select( {particle_idx}, {particle_count} ).read( particles.color2 ); pdata.getDataSet("px").select( {particle_idx}, {particle_count} ).read( particles.px ); pdata.getDataSet("py").select( {particle_idx}, {particle_count} ).read( particles.py ); pdata.getDataSet("pz").select( {particle_idx}, {particle_count} ).read( particles.pz ); pdata.getDataSet("e").select( {particle_idx}, {particle_count} ).read( particles.e ); pdata.getDataSet("m").select( {particle_idx}, {particle_count} ).read( particles.m ); pdata.getDataSet("lifetime").select( {particle_idx}, {particle_count} ).read( particles.lifetime ); pdata.getDataSet("spin").select( {particle_idx}, {particle_count} ).read( particles.spin ); event_idx += count; particle_idx += particle_count; return count; } }; HDF5Reader::HDF5Reader(std::string const & filename): impl_{ new HDF5Reader::HDF5ReaderImpl{filename}, HDF5Reader::HDF5ReaderImplDeleter{} } {} bool HDF5Reader::read_event() { if(impl_->cur_event == cend(impl_->records)) { // end of active chunk, read new events from file const auto events_read = impl_->read_event_records(chunk_size); impl_->cur_event = cbegin(impl_->records); if(events_read == 0) return false; } impl_->hepeup = *impl_->cur_event; ++impl_->cur_event; return true; } namespace { static const std::string nothing = ""; } std::string const & HDF5Reader::header() const { return nothing; } LHEF::HEPRUP const & HDF5Reader::heprup() const { return impl_->heprup; } LHEF::HEPEUP const & HDF5Reader::hepeup() const { return impl_->hepeup; } HEJ::optional<int> HDF5Reader::number_events() const { return impl_->nevents; } } #else // no HDF5 support namespace HEJ { class HDF5Reader::HDF5ReaderImpl{}; HDF5Reader::HDF5Reader(std::string const &){ throw std::invalid_argument{ "Failed to create HDF5 reader: " "HEJ 2 was built without HDF5 support" }; } bool HDF5Reader::read_event() { throw std::logic_error{"unreachable"}; } std::string const & HDF5Reader::header() const { throw std::logic_error{"unreachable"}; } LHEF::HEPRUP const & HDF5Reader::heprup() const { throw std::logic_error{"unreachable"}; } LHEF::HEPEUP const & HDF5Reader::hepeup() const { throw std::logic_error{"unreachable"}; } HEJ::optional<int> HDF5Reader::number_events() const { throw std::logic_error{"unreachable"}; } } #endif namespace HEJ { void HDF5Reader::HDF5ReaderImplDeleter::operator()(HDF5ReaderImpl* p) { delete p; } }