GenEvent.h

// -*- C++ -*-
//
// This file is part of HepMC
// Copyright (C) 2014-2023 The HepMC collaboration (see AUTHORS for details)
//
#ifndef HEPMC3_GENEVENT_H
#define HEPMC3_GENEVENT_H
 
#include "HepMC3/Units.h"
#include "HepMC3/GenParticle_fwd.h"
#include "HepMC3/GenVertex_fwd.h"
#include "HepMC3/GenPdfInfo_fwd.h"
#include "HepMC3/GenHeavyIon_fwd.h"
#include "HepMC3/GenCrossSection_fwd.h"
 
#if !defined(__CINT__)
#include "HepMC3/GenHeavyIon.h"
#include "HepMC3/GenPdfInfo.h"
#include "HepMC3/GenCrossSection.h"
#include "HepMC3/GenRunInfo.h"
#include <mutex>
#endif // __CINT__
 
#ifdef HEPMC3_ROOTIO
class TBuffer;
#endif
 
#ifdef HEPMC3_PROTOBUFIO
namespace HepMC3_pb {
class GenEventData;
}
#endif
 
 
namespace HepMC3 {
 
struct GenEventData;

class GenEvent {
 
public:

    GenEvent(Units::MomentumUnit mu = Units::GEV,
             Units::LengthUnit lu = Units::MM);
 
#if !defined(__CINT__)

    GenEvent(std::shared_ptr<GenRunInfo> run,
             Units::MomentumUnit mu = Units::GEV,
             Units::LengthUnit lu = Units::MM);

    GenEvent(const GenEvent&);

    ~GenEvent();

    GenEvent& operator=(const GenEvent&);


    const std::vector<ConstGenParticlePtr>& particles() const;
    const std::vector<ConstGenVertexPtr>& vertices() const;

    const std::vector<GenParticlePtr>& particles() { return m_particles; }
    const std::vector<GenVertexPtr>& vertices() { return m_vertices; }

 


    inline int particles_size() const { return m_particles.size(); }
    inline bool particles_empty() const { return m_particles.empty(); }
    inline int vertices_size() const { return m_vertices.size(); }
    inline bool vertices_empty() const { return m_vertices.empty(); }
 


    const std::vector<double>& weights() const { return m_weights; }
    std::vector<double>& weights() { return m_weights; }
    double weight(const unsigned long& index=0) const { if ( index < weights().size() ) return weights().at(index); else  throw std::runtime_error("GenEvent::weight(const unsigned long&): weight index outside of range"); return 0.0; }
    double weight(const std::string& name) const {
        if (!run_info()) throw std::runtime_error("GenEvent::weight(const std::string&): named access to event weights requires the event to have a GenRunInfo");
        return weight(run_info()->weight_index(name));
    }

    double& weight(const std::string& name) {
        if (!run_info()) throw std::runtime_error("GenEvent::weight(const std::string&): named access to event weights requires the event to have a GenRunInfo");
        int pos = run_info()->weight_index(name);
        if ( pos < 0 ) throw std::runtime_error("GenEvent::weight(const std::string&): no weight with given name in this run");
        if ( pos >= int(m_weights.size())) throw std::runtime_error("GenEvent::weight(const std::string&): weight index outside of range");
        return m_weights[pos];
    }

    const std::vector<std::string>& weight_names() const {
        if (!run_info()) throw std::runtime_error("GenEvent::weight_names(): access to event weight names requires the event to have a GenRunInfo");
        const std::vector<std::string>& weightnames = run_info()->weight_names();
        if (weightnames.empty()) throw std::runtime_error("GenEvent::weight_names(): no event weight names are registered for this run");
        return weightnames;
    }

 


    std::shared_ptr<GenRunInfo> run_info() const {
        return m_run_info;
    }

    void set_run_info(std::shared_ptr<GenRunInfo> run) {
        m_run_info = run;
        if ( run && !run->weight_names().empty() )
            m_weights.resize(run->weight_names().size(), 1.0);
    }

    int  event_number() const { return m_event_number; }
    void set_event_number(const int& num) { m_event_number = num; }

    const Units::MomentumUnit& momentum_unit() const { return m_momentum_unit; }
    const Units::LengthUnit& length_unit() const { return m_length_unit; }
    void set_units( Units::MomentumUnit new_momentum_unit, Units::LengthUnit new_length_unit);

    GenHeavyIonPtr heavy_ion() { return attribute<GenHeavyIon>("GenHeavyIon"); }
    ConstGenHeavyIonPtr heavy_ion() const { return attribute<GenHeavyIon>("GenHeavyIon"); }
    void set_heavy_ion(GenHeavyIonPtr hi) { add_attribute("GenHeavyIon",hi); }

    GenPdfInfoPtr pdf_info() { return attribute<GenPdfInfo>("GenPdfInfo"); }
    ConstGenPdfInfoPtr pdf_info() const { return attribute<GenPdfInfo>("GenPdfInfo"); }
    void set_pdf_info(GenPdfInfoPtr pi) { add_attribute("GenPdfInfo",pi); }

    GenCrossSectionPtr cross_section() { return attribute<GenCrossSection>("GenCrossSection"); }
    ConstGenCrossSectionPtr cross_section() const { return attribute<GenCrossSection>("GenCrossSection"); }
    void set_cross_section(GenCrossSectionPtr cs) { add_attribute("GenCrossSection",cs); }

 


    const FourVector& event_pos() const;

    std::vector<ConstGenParticlePtr> beams() const;

    std::vector<ConstGenParticlePtr> beams(const int status) const;
 

    const std::vector<GenParticlePtr> & beams();

    void shift_position_by( const FourVector & delta );

    void shift_position_to( const FourVector & newpos ) {
        const FourVector delta = newpos - event_pos();
        shift_position_by(delta);
    }

    bool boost( const FourVector&  delta );
    bool rotate( const FourVector&  delta );
    bool reflect(const int axis);

 

    void add_attribute(const std::string &name, const std::shared_ptr<Attribute> &att,  const int& id = 0);

    void add_attributes(const std::vector<std::string> &names, const std::vector<std::shared_ptr<Attribute> > &atts,  const std::vector<int>& ids);

    void add_attributes(const std::string& name, const std::vector<std::shared_ptr<Attribute> > &atts,  const std::vector<int>& ids);

    void add_attributes(const std::string& name, const std::vector<std::pair<int, std::shared_ptr<Attribute> > > &atts);
 

    void remove_attribute(const std::string &name,  const int& id = 0);

    template<class T>
    std::shared_ptr<T> attribute(const std::string &name,  const int& id = 0) const;

    std::string attribute_as_string(const std::string &name,  const int& id = 0) const;

    std::vector<std::string> attribute_names( const int& id = 0) const;

    std::map< std::string, std::map<int, std::shared_ptr<Attribute> > > attributes() const {
        std::lock_guard<std::recursive_mutex> lock(m_lock_attributes);
        return m_attributes;
    }

 


    void add_particle( GenParticlePtr p );

    void add_vertex( GenVertexPtr v );

    void remove_particle( GenParticlePtr p );

    void remove_particles( std::vector<GenParticlePtr> v );

    void remove_vertex( GenVertexPtr v );

    void add_tree( const std::vector<GenParticlePtr> &parts );

    void reserve(const size_t& parts, const size_t& verts = 0);

    void clear();



    void set_beam_particles(GenParticlePtr p1, GenParticlePtr p2);
 

 
    void add_beam_particle(GenParticlePtr p1);
 

 
#endif // __CINT__
 


    void write_data(GenEventData &data) const;

    void read_data(const GenEventData &data);
 
#ifdef HEPMC3_ROOTIO
    void Streamer(TBuffer &b);
#endif
 
#ifdef HEPMC3_PROTOBUFIO
    void read_data(HepMC3_pb::GenEventData const &data);
#endif
 
private:

 
#if !defined(__CINT__)

    std::vector<GenParticlePtr> m_particles;
    std::vector<GenVertexPtr> m_vertices;

    int m_event_number = 0;

    std::vector<double> m_weights;

    Units::MomentumUnit m_momentum_unit = Units::GEV;
    Units::LengthUnit m_length_unit = Units::MM;

    GenVertexPtr m_rootvertex;

    std::shared_ptr<GenRunInfo> m_run_info;

    mutable std::map< std::string, std::map<int, std::shared_ptr<Attribute> > > m_attributes;

    typedef std::map< std::string, std::map<int, std::shared_ptr<Attribute> > >::value_type att_key_t;

    typedef std::map<int, std::shared_ptr<Attribute> >::value_type att_val_t;

    mutable std::recursive_mutex m_lock_attributes;
#endif // __CINT__

 
};
 
 
 
#if !defined(__CINT__)
//
// Template methods
//
template<class T>
std::shared_ptr<T> GenEvent::attribute(const std::string &name,  const int& id) const {
    std::lock_guard<std::recursive_mutex> lock(m_lock_attributes);
    std::map< std::string, std::map<int, std::shared_ptr<Attribute> > >::iterator i1 = m_attributes.find(name);
    if ( i1 == m_attributes.end() ) {
        if ( id == 0 && run_info() ) {
            return run_info()->attribute<T>(name);
        }
        return std::shared_ptr<T>();
    }
 
    std::map<int, std::shared_ptr<Attribute> >::iterator i2 = i1->second.find(id);
    if (i2 == i1->second.end() ) return std::shared_ptr<T>();
 
    if (!i2->second->is_parsed() ) {
 
        std::shared_ptr<T> att = std::make_shared<T>();
        att->m_event = this;
 
        if ( id > 0 && id <= int(particles().size()) ) {
            att->m_particle = m_particles[id - 1];
        }
        if ( id < 0 && -id <= int(vertices().size()) ) {
            att->m_vertex = m_vertices[-id - 1];
        }
        if ( att->from_string(i2->second->unparsed_string()) &&
                att->init() ) {
            // update map with new pointer
            i2->second = att;
            return att;
        } else {
            return std::shared_ptr<T>();
        }
    }
    else return std::dynamic_pointer_cast<T>(i2->second);
}
#endif // __CINT__
 
 
} // namespace HepMC3
#endif