diff --git a/src/HDF5Writer.cc b/src/HDF5Writer.cc index 0e0b33b..ea90b58 100644 --- a/src/HDF5Writer.cc +++ b/src/HDF5Writer.cc @@ -1,432 +1,435 @@ /** * \authors The HEJ collaboration (see AUTHORS for details) * \date 2019-2020 * \copyright GPLv2 or later */ #include "HEJ/HDF5Writer.hh" #include #include "LHEF/LHEF.h" #include "HEJ/ConfigFlags.hh" #ifdef HEJ_BUILD_WITH_HDF5 #include #include #include #include #include #include #include "highfive/H5File.hpp" #include "HEJ/Event.hh" #include "HEJ/event_types.hh" #else #include "HEJ/exceptions.hh" #endif #ifdef HEJ_BUILD_WITH_HDF5 namespace HEJ { using HighFive::File; using HighFive::DataSpace; namespace { using std::size_t; constexpr size_t CHUNK_SIZE = 1000; constexpr unsigned COMPRESSION_LEVEL = 3; size_t to_index(event_type::EventType const type){ return type==0?0:std::floor(std::log2(static_cast(type)))+1; } template void write_dataset(HighFive::Group & group, std::string const & name, T val) { using data_t = std::decay_t; group.createDataSet(name, DataSpace::From(val)).write(val); } template void write_dataset( HighFive::Group & group, std::string const & name, std::vector const & val ) { using data_t = std::decay_t; group.createDataSet(name, DataSpace::From(val)).write(val); } struct Cache { explicit Cache(size_t capacity): capacity{capacity} { nparticles.reserve(capacity); start.reserve(capacity); pid.reserve(capacity); weight.reserve(capacity); scale.reserve(capacity); fscale.reserve(capacity); rscale.reserve(capacity); aqed.reserve(capacity); aqcd.reserve(capacity); trials.reserve(capacity); npLO.reserve(capacity); npNLO.reserve(capacity); } void fill(Event const & ev) { const auto hepeup = to_HEPEUP(ev, nullptr); // HEJ event to get nice wrapper const auto num_partons = std::distance(ev.cbegin_partons(), ev.cend_partons()); assert(num_partons>0); // Number of partons for LO matching, HEJ requires at least 2 partons npLO.emplace_back(num_partons>1?num_partons-2:num_partons); // Number of real emissions in NLO, HEJ is LO -> -1 npNLO.emplace_back(-1); fill_event_params(hepeup); fill_event_particles(hepeup); } void fill_event_params(LHEF::HEPEUP const & ev) { nparticles.emplace_back(ev.NUP); start.emplace_back(particle_pos); pid.emplace_back(ev.IDPRUP); weight.emplace_back(ev.XWGTUP); scale.emplace_back(ev.SCALUP); fscale.emplace_back(ev.scales.muf); rscale.emplace_back(ev.scales.mur); aqed.emplace_back(ev.AQEDUP); aqcd.emplace_back(ev.AQCDUP); // set first trial=1 for first event // -> sum(trials) = 1 -> xs=sum(weights)/sum(trials) as in Sherpa if(particle_pos == 0){ trials.emplace_back(1.); } else { trials.emplace_back(0.); } particle_pos += ev.NUP; } void fill_event_particles(LHEF::HEPEUP const & ev) { id.insert(end(id), begin(ev.IDUP), end(ev.IDUP)); status.insert(end(status), begin(ev.ISTUP), end(ev.ISTUP)); lifetime.insert(end(lifetime), begin(ev.VTIMUP), end(ev.VTIMUP)); spin.insert(end(spin), begin(ev.SPINUP), end(ev.SPINUP)); for(int i = 0; i < ev.NUP; ++i) { mother1.emplace_back(ev.MOTHUP[i].first); mother2.emplace_back(ev.MOTHUP[i].second); color1.emplace_back(ev.ICOLUP[i].first); color2.emplace_back(ev.ICOLUP[i].second); px.emplace_back(ev.PUP[i][0]); py.emplace_back(ev.PUP[i][1]); pz.emplace_back(ev.PUP[i][2]); e.emplace_back(ev.PUP[i][3]); m.emplace_back(ev.PUP[i][4]); } } bool is_full() const { return nparticles.size() >= capacity; } void clear() { nparticles.clear(); start.clear(); pid.clear(); id.clear(); status.clear(); mother1.clear(); mother2.clear(); color1.clear(); color2.clear(); weight.clear(); scale.clear(); fscale.clear(); rscale.clear(); aqed.clear(); aqcd.clear(); trials.clear(); npLO.clear(); npNLO.clear(); px.clear(); py.clear(); pz.clear(); e.clear(); m.clear(); lifetime.clear(); spin.clear(); } size_t capacity; std::vector nparticles, start, pid, id, status, mother1, mother2, color1, color2, npLO, npNLO; std::vector weight, scale, fscale, rscale, aqed, aqcd, trials, px, py, pz, e, m, lifetime, spin; size_t particle_pos = 0; }; } // namespace struct HDF5Writer::HDF5WriterImpl: EventWriter{ File file; LHEF::HEPRUP heprup; Cache cache{CHUNK_SIZE}; size_t event_idx = 0; size_t particle_idx = 0; HDF5WriterImpl(std::string const & filename, LHEF::HEPRUP && hepr): file{filename, File::ReadWrite | File::Create | File::Truncate}, heprup{std::forward(hepr)} { // TODO: code duplication with Les Houches Writer const size_t max_number_types = to_index(event_type::last_type)+1; heprup.NPRUP = max_number_types; // ids of event types heprup.LPRUP.clear(); heprup.LPRUP.reserve(max_number_types); heprup.LPRUP.emplace_back(0); for(size_t i=event_type::first_type+1; i<=event_type::last_type; i*=2) { heprup.LPRUP.emplace_back(i); } heprup.XSECUP = std::vector(max_number_types, 0.); heprup.XERRUP = std::vector(max_number_types, 0.); heprup.XMAXUP = std::vector(max_number_types, 0.); write_init(); create_event_group(); create_particle_group(); } void write_init() { auto init = file.createGroup("init"); write_dataset(init, "PDFgroupA" , heprup.PDFGUP.first); write_dataset(init, "PDFgroupB" , heprup.PDFGUP.second); write_dataset(init, "PDFsetA" , heprup.PDFSUP.first); write_dataset(init, "PDFsetB" , heprup.PDFSUP.second); write_dataset(init, "beamA" , heprup.IDBMUP.first); write_dataset(init, "beamB" , heprup.IDBMUP.second); write_dataset(init, "energyA" , heprup.EBMUP.first); write_dataset(init, "energyB" , heprup.EBMUP.second); write_dataset(init, "numProcesses" , heprup.NPRUP); write_dataset(init, "weightingStrategy", heprup.IDWTUP); auto proc_info = file.createGroup("procInfo"); write_dataset(proc_info, "procId", heprup.LPRUP); } static HighFive::DataSetCreateProps const & hdf5_chunk() { static const auto props = [](){ HighFive::DataSetCreateProps props; props.add(HighFive::Chunking({CHUNK_SIZE})); props.add(HighFive::Deflate(COMPRESSION_LEVEL)); return props; }(); return props; } void create_event_group() { static const auto dim = DataSpace({0}, {DataSpace::UNLIMITED}); auto events = file.createGroup("event"); events.createDataSet("nparticles", dim, hdf5_chunk()); events.createDataSet("start", dim, hdf5_chunk()); events.createDataSet("pid", dim, hdf5_chunk()); events.createDataSet("weight", dim, hdf5_chunk()); events.createDataSet("scale", dim, hdf5_chunk()); events.createDataSet("fscale", dim, hdf5_chunk()); events.createDataSet("rscale", dim, hdf5_chunk()); events.createDataSet("aqed", dim, hdf5_chunk()); events.createDataSet("aqcd", dim, hdf5_chunk()); events.createDataSet("trials", dim, hdf5_chunk()); events.createDataSet("npLO", dim, hdf5_chunk()); events.createDataSet("npNLO", dim, hdf5_chunk()); } void resize_event_group(size_t new_size) { auto events = file.getGroup("event"); events.getDataSet("nparticles").resize({new_size}); events.getDataSet("start").resize({new_size}); events.getDataSet("pid").resize({new_size}); events.getDataSet("weight").resize({new_size}); events.getDataSet("scale").resize({new_size}); events.getDataSet("fscale").resize({new_size}); events.getDataSet("rscale").resize({new_size}); events.getDataSet("aqed").resize({new_size}); events.getDataSet("aqcd").resize({new_size}); events.getDataSet("trials").resize({new_size}); events.getDataSet("npLO").resize({new_size}); events.getDataSet("npNLO").resize({new_size}); } void create_particle_group() { static const auto dim = DataSpace({0}, {DataSpace::UNLIMITED}); auto particles = file.createGroup("particle"); particles.createDataSet("id", dim, hdf5_chunk()); particles.createDataSet("status", dim, hdf5_chunk()); particles.createDataSet("mother1", dim, hdf5_chunk()); particles.createDataSet("mother2", dim, hdf5_chunk()); particles.createDataSet("color1", dim, hdf5_chunk()); particles.createDataSet("color2", dim, hdf5_chunk()); particles.createDataSet("px", dim, hdf5_chunk()); particles.createDataSet("py", dim, hdf5_chunk()); particles.createDataSet("pz", dim, hdf5_chunk()); particles.createDataSet("e", dim, hdf5_chunk()); particles.createDataSet("m", dim, hdf5_chunk()); particles.createDataSet("lifetime", dim, hdf5_chunk()); particles.createDataSet("spin", dim, hdf5_chunk()); } void resize_particle_group(size_t new_size) { auto particles = file.getGroup("particle"); particles.getDataSet("id").resize({new_size}); particles.getDataSet("status").resize({new_size}); particles.getDataSet("mother1").resize({new_size}); particles.getDataSet("mother2").resize({new_size}); particles.getDataSet("color1").resize({new_size}); particles.getDataSet("color2").resize({new_size}); particles.getDataSet("px").resize({new_size}); particles.getDataSet("py").resize({new_size}); particles.getDataSet("pz").resize({new_size}); particles.getDataSet("e").resize({new_size}); particles.getDataSet("m").resize({new_size}); particles.getDataSet("lifetime").resize({new_size}); particles.getDataSet("spin").resize({new_size}); } void write(Event const & ev) override { cache.fill(ev); if(cache.is_full()) { dump_cache(); } const double wt = ev.central().weight; const size_t idx = to_index(ev.type()); heprup.XSECUP[idx] += wt; heprup.XERRUP[idx] += wt*wt; if(wt > heprup.XMAXUP[idx]){ heprup.XMAXUP[idx] = wt; } } void dump_cache() { write_event_params(); write_event_particles(); cache.clear(); } void write_event_params() { auto events = file.getGroup("event"); // choose arbitrary dataset to find size const auto dataset = events.getDataSet("nparticles"); const size_t size = dataset.getSpace().getDimensions().front(); resize_event_group(size + cache.nparticles.size()); // NOLINTNEXTLINE #define WRITE_FROM_CACHE(GROUP, PROPERTY) \ GROUP.getDataSet(#PROPERTY).select({size}, {cache.PROPERTY.size()}).write(cache.PROPERTY) WRITE_FROM_CACHE(events, nparticles); WRITE_FROM_CACHE(events, start); WRITE_FROM_CACHE(events, pid); WRITE_FROM_CACHE(events, weight); WRITE_FROM_CACHE(events, scale); WRITE_FROM_CACHE(events, fscale); WRITE_FROM_CACHE(events, rscale); WRITE_FROM_CACHE(events, aqed); WRITE_FROM_CACHE(events, aqcd); WRITE_FROM_CACHE(events, trials); WRITE_FROM_CACHE(events, npLO); WRITE_FROM_CACHE(events, npNLO); } void write_event_particles() { auto particles = file.getGroup("particle"); // choose arbitrary dataset to find size const auto dataset = particles.getDataSet("id"); const size_t size = dataset.getSpace().getDimensions().front(); resize_particle_group(size + cache.id.size()); WRITE_FROM_CACHE(particles, id); WRITE_FROM_CACHE(particles, status); WRITE_FROM_CACHE(particles, lifetime); WRITE_FROM_CACHE(particles, spin); WRITE_FROM_CACHE(particles, mother1); WRITE_FROM_CACHE(particles, mother2); WRITE_FROM_CACHE(particles, color1); WRITE_FROM_CACHE(particles, color2); WRITE_FROM_CACHE(particles, px); WRITE_FROM_CACHE(particles, py); WRITE_FROM_CACHE(particles, pz); WRITE_FROM_CACHE(particles, e); WRITE_FROM_CACHE(particles, m); } #undef WRITE_FROM_CACHE void finish() override { if(finished()) throw std::ios_base::failure("HDF5Writer writer already finished."); EventWriter::finish(); dump_cache(); auto proc_info = file.getGroup("procInfo"); write_dataset(proc_info, "xSection", heprup.XSECUP); + for(auto & xs_err: heprup.XERRUP) { + xs_err = std::sqrt(xs_err); + } write_dataset(proc_info, "error", heprup.XERRUP); write_dataset(proc_info, "unitWeight", heprup.XMAXUP); file.flush(); } ~HDF5WriterImpl() override { finish_or_abort(this, "HDF5Writer"); } }; HDF5Writer::HDF5Writer(std::string const & filename, LHEF::HEPRUP heprup): impl_{std::make_unique( filename, std::move(heprup))} {} void HDF5Writer::write(Event const & ev){ impl_->write(ev); } void HDF5Writer::finish(){ impl_->finish(); } } // namespace HEJ #else // no HDF5 support namespace HEJ { struct HDF5Writer::HDF5WriterImpl{}; HDF5Writer::HDF5Writer(std::string const & /*file*/, LHEF::HEPRUP /*heprup*/){ throw std::invalid_argument{ "Failed to create HDF5 writer: " "HEJ 2 was built without HDF5 support" }; } void HDF5Writer::write(Event const & /*ev*/){ assert(false); } void HDF5Writer::finish(){ assert(false); } } #endif namespace HEJ { HDF5Writer::~HDF5Writer() = default; }