Merge branch 'master' of github.com:cepc/CEPCSW

Generator/CMakeLists.txt

@@ -15,7 +15,9 @@ gaudi_add_module(GenAlgo
                          src/SLCIORdr.cpp
                          src/HepMCRdr.cpp
                          src/GtGunTool.cpp
-
+                         # ------- Beam Background -------
+                         src/GtBeamBackgroundTool.cpp
+                         src/BeamBackgroundFileParserV0.cpp
                  LINK ${ROOT_LIBRARIES}
                       k4FWCore::k4FWCore 
                       Gaudi::GaudiAlgLib

Generator/src/BeamBackgroundFileParserV0.cpp

+
+#include "BeamBackgroundFileParserV0.h"
+#include <sstream>
+#include <cmath>
+
+BeamBackgroundFileParserV0::BeamBackgroundFileParserV0(const std::string& filename,
+                                                       int pdgid,
+                                                       double beam_energy) {
+    m_input.open(filename.c_str());
+    m_pdgid = pdgid;
+    m_beam_energy = beam_energy;
+}
+
+bool BeamBackgroundFileParserV0::load(IBeamBackgroundFileParser::BeamBackgroundData& result) {
+
+    if (not m_input.good()) {
+        return false;
+    }
+
+    // read one record
+    std::string tmpline;
+    // the format
+    double generation_point;
+    int loss_turn;
+    double z; // unit: m
+    double x; // unit: m
+    double y; // unit: m
+    double cosx; // 
+    double cosy; //
+    double dz; // unit: m
+    double dp; // unit: relative to the E
+
+    while(m_input.good()) {
+        std::getline(m_input, tmpline);
+        std::stringstream ss;
+        ss << tmpline;
+        ss >> generation_point; if (ss.fail()) { continue; }
+        ss >> loss_turn;        if (ss.fail()) { continue; }
+        ss >> z;                if (ss.fail()) { continue; }
+        ss >> x;                if (ss.fail()) { continue; }
+        ss >> y;                if (ss.fail()) { continue; }
+        ss >> cosx;             if (ss.fail()) { continue; }
+        ss >> cosy;             if (ss.fail()) { continue; }
+        ss >> dz;               if (ss.fail()) { continue; }
+        ss >> dp;               if (ss.fail()) { continue; }
+
+        double p = m_beam_energy*(1+dp);
+
+        // Now, we get a almost valid data
+        const double m2mm = 1e3; // convert from m to mm
+        result.pdgid = m_pdgid;
+        result.x     = x * m2mm; 
+        result.y     = y * m2mm; 
+        result.z     = (z+dz) * m2mm;
+
+        result.px    = p * cosx;
+        result.py    = p * cosy;
+        result.pz    = p * std::sqrt(1-cosx*cosx-cosy*cosy);
+
+        result.mass  = 0.000511; // assume e-/e+, mass is 0.511 MeV
+
+        return true;
+    }
+    return false;
+}

Generator/src/BeamBackgroundFileParserV0.h

+#ifndef BeamBackgroundFileParserV0_h
+#define BeamBackgroundFileParserV0_h
+
+#include "IBeamBackgroundFileParser.h"
+
+#include <fstream>
+
+class BeamBackgroundFileParserV0: public IBeamBackgroundFileParser {
+public:
+    BeamBackgroundFileParserV0(const std::string& filename, int pdgid, double beam_energy);
+
+    bool load(IBeamBackgroundFileParser::BeamBackgroundData&);
+private:
+    std::ifstream m_input;
+
+    int m_pdgid;
+    double m_beam_energy;
+};
+
+#endif

Generator/src/GtBeamBackgroundTool.cpp

+#include "GtBeamBackgroundTool.h"
+#include "IBeamBackgroundFileParser.h"
+#include "BeamBackgroundFileParserV0.h"
+
+#include "TVector3.h" // for rotation
+DECLARE_COMPONENT(GtBeamBackgroundTool)
+
+StatusCode GtBeamBackgroundTool::initialize() {
+    // check the consistency of the properties
+
+    // create the instances of the background parsers
+
+    for (auto& [label, inputfn]: m_inputmaps) {
+        double beamE = 120.;
+        auto itBeamE = m_Ebeammaps.find(label);
+        if (itBeamE != m_Ebeammaps.end()) {
+            beamE = itBeamE->second;
+        }
+        info() << "Initializing beam background ... "
+               << label << " "
+               << beamE << " "
+               << inputfn
+               << endmsg;
+        m_beaminputs[label] = std::make_shared<BeamBackgroundFileParserV0>(inputfn, 11, beamE);
+    }
+
+    // check the size
+    if (m_beaminputs.empty()) {
+        error() << "Empty Beam Background File Parser. " << endmsg;
+        return StatusCode::FAILURE;
+    }
+
+    return StatusCode::SUCCESS;
+}
+
+StatusCode GtBeamBackgroundTool::finalize() {
+    return StatusCode::SUCCESS;
+}
+
+
+bool GtBeamBackgroundTool::mutate(MyHepMC::GenEvent& event) {
+    if (m_beaminputs.empty()) {
+        return false;
+    }
+    // TODO: should sample according to the rates
+    // dummy: get one and stop to generate
+    for (auto& [label, parser]: m_beaminputs) {
+        IBeamBackgroundFileParser::BeamBackgroundData beamdata;
+        auto isok = parser->load(beamdata);
+        if (not isok) {
+            error() << "Failed to load beam background data from the parser " << label << endmsg;
+            return false;
+        }
+        // fill the value
+        float charge = beamdata.pdgid == 11 ? -1: 1;
+
+        TVector3 pos(beamdata.x,beamdata.y,beamdata.z);
+        TVector3 mom(beamdata.px,beamdata.py,beamdata.pz);
+
+        auto itrot = m_rotYmaps.find(label);
+        if (itrot  != m_rotYmaps.end() ) {
+            info() << "Apply rotation along Y " << itrot->second << endmsg;
+
+            pos.RotateY(itrot->second);
+            mom.RotateY(itrot->second);
+        }
+
+        // create the MC particle
+        edm4hep::MCParticle mcp = event.m_mc_vec.create();
+        mcp.setPDG(beamdata.pdgid);
+        mcp.setGeneratorStatus(1);
+        mcp.setSimulatorStatus(1);
+        mcp.setCharge(static_cast<float>(charge));
+        mcp.setTime(beamdata.t);
+        mcp.setMass(beamdata.mass);
+        mcp.setVertex(edm4hep::Vector3d(pos.X(), pos.Y(), pos.Z())); 
+        mcp.setMomentum(edm4hep::Vector3f(mom.X(), mom.Y(), mom.Z()));
+
+    }
+    
+    return true;
+}
+
+bool GtBeamBackgroundTool::finish() {
+    return true;
+}
+
+bool GtBeamBackgroundTool::configure_gentool() {
+
+    return true;
+}

Generator/src/GtBeamBackgroundTool.h

+#ifndef GtBeamBackgroundTool_h
+#define GtBeamBackgroundTool_h
+
+/*
+ * Description:
+ *   This tool is used to simulation the non-collision beam backgrounds.
+ *
+ *   The properties:
+ *     - InputFileMap
+ *         this is a map to store the label and the input filename
+ *     - InputFormatMap
+ *         this is a map to store the label and the input format
+ *     - InputRateMap
+ *         this is a map to store the label and the rate
+ * 
+ *     Note: the label (key) should be consistent
+ *
+ * About the design:
+ *   IBeamBackgroundFileParser is the interface to load the next event.
+ *   Different file formats should be implemented in the corresponding parsers. 
+ *   The format will be used to create the corresponding instance.
+ *
+ * Author:
+ *   Tao Lin <lintao AT ihep.ac.cn>
+ */
+
+#include <GaudiKernel/AlgTool.h>
+#include <Gaudi/Property.h>
+#include "IGenTool.h"
+#include "IBeamBackgroundFileParser.h"
+
+#include <vector>
+#include <map>
+
+
+class GtBeamBackgroundTool: public extends<AlgTool, IGenTool> {
+public:
+    using extends::extends;
+
+    // Overriding initialize and finalize
+    StatusCode initialize() override;
+    StatusCode finalize() override;
+
+    // IGenTool
+    bool mutate(MyHepMC::GenEvent& event) override;
+    bool finish() override;
+    bool configure_gentool() override;
+
+private:
+    Gaudi::Property<std::map<std::string, std::string>> m_inputmaps{this, "InputFileMap"};
+    Gaudi::Property<std::map<std::string, std::string>> m_fomatmaps{this, "InputFormatMap"};
+    Gaudi::Property<std::map<std::string, double>>      m_ratemaps {this, "InputRateMap"};
+
+    // unit of beam energy: GeV
+    Gaudi::Property<std::map<std::string, double>>      m_Ebeammaps{this, "InputBeamEnergyMap"};
+
+    // unit of the rotation along Y: rad
+    Gaudi::Property<std::map<std::string, double>>      m_rotYmaps {this, "RotationAlongYMap"};
+
+private:
+    std::map<std::string, std::shared_ptr<IBeamBackgroundFileParser>> m_beaminputs;
+
+};
+
+#endif

Generator/src/IBeamBackgroundFileParser.h

+#ifndef IBeamBackgroundFileParser_h
+#define IBeamBackgroundFileParser_h
+
+/*
+ * Description:
+ *   This interface is used to load the beam background information, such as:
+ *     - pdgid (optional)
+ *         About the pdgid, it will be e+/e- in most cases.
+ *     - x/y/z
+ *     - t (optional) 
+ *     - px/py/pz
+ *         About the time, it could be set in the GtBeamBackgroundTool.
+ *
+ * Author:
+ *   Tao Lin <lintao AT ihep.ac.cn>
+ */
+
+class IBeamBackgroundFileParser {
+public:
+    // Internal used Data
+    struct BeamBackgroundData {
+        int pdgid;
+
+        double x; // unit: mm
+        double y; // unit: mm
+        double z; // unit: mm
+        double t; // unit: ns
+
+        double px; // unit: GeV
+        double py; // unit: GeV
+        double pz; // unit: GeV
+        double mass; // unit: GeV
+
+        BeamBackgroundData() 
+          : pdgid(11), x(0), y(0), z(0), t(0), 
+            px(0), py(0), pz(0), mass(0) {}
+        
+    };
+
+    // return false if failed to load the data
+    virtual bool load(BeamBackgroundData&) = 0;
+};
+
+#endif