From 63eaa7f40f03710371aed61baa45c194f2a957e1 Mon Sep 17 00:00:00 2001
From: lintao <lintao51@gmail.com>
Date: Tue, 29 Nov 2022 17:09:34 +0800
Subject: [PATCH] WIP: add R/Z min/max in gun.

---
 Generator/src/GtGunTool.cpp | 64 +++++++++++++++++++++++++++++++++++++
 Generator/src/GtGunTool.h   |  9 ++++++
 2 files changed, 73 insertions(+)

diff --git a/Generator/src/GtGunTool.cpp b/Generator/src/GtGunTool.cpp
index af2a98dd..7f9b56a3 100644
--- a/Generator/src/GtGunTool.cpp
+++ b/Generator/src/GtGunTool.cpp
@@ -32,6 +32,39 @@ GtGunTool::initialize() {
         error() << "Mismatched PositionZs and particles." << endmsg;
         return StatusCode::FAILURE;
     }
+
+
+    if (m_posZmins.value().size()
+        && m_posZmins.value().size() != m_particles.value().size()) {
+        error() << "Mismatched PosZmins and particles." << endmsg;
+        return StatusCode::FAILURE;
+    }
+    if (m_posZmaxs.value().size()
+        && m_posZmaxs.value().size() != m_particles.value().size()) {
+        error() << "Mismatched PosZmaxs and particles." << endmsg;
+        return StatusCode::FAILURE;
+    }
+
+    if (m_posZmins.value().size() != m_posZmaxs.value().size()) {
+        error() << "Mismatched PosZmins and PosZmaxs." << endmsg;
+        return StatusCode::FAILURE;
+    }
+
+    if (m_posRmins.value().size()
+        && m_posRmins.value().size() != m_particles.value().size()) {
+        error() << "Mismatched PosRmins and particles." << endmsg;
+        return StatusCode::FAILURE;
+    }
+    if (m_posRmaxs.value().size()
+        && m_posRmaxs.value().size() != m_particles.value().size()) {
+        error() << "Mismatched PosRmaxs and particles." << endmsg;
+        return StatusCode::FAILURE;
+    }
+
+    if (m_posRmins.value().size() != m_posRmaxs.value().size()) {
+        error() << "Mismatched PosRmins and PosRmaxs." << endmsg;
+        return StatusCode::FAILURE;
+    }
     
     // Energy
     if (m_energymins.value().size() != m_particles.value().size()) {
@@ -137,10 +170,41 @@ GtGunTool::mutate(MyHepMC::GenEvent& event) {
         double x = 0;
         double y = 0;
         double z = 0;
+
+        // ==================================
+        // 1. if there is fixed positions
+        // ==================================
         if (i<m_positionXs.value().size()) { x = m_positionXs.value()[i]; }
         if (i<m_positionYs.value().size()) { y = m_positionYs.value()[i]; }
         if (i<m_positionZs.value().size()) { z = m_positionZs.value()[i]; }
 
+        // ==================================
+        // 2. if there are varied positions
+        // ==================================
+        if (i<m_posZmins.value().size() and i<m_posZmaxs.value().size()) {
+            double zmin = m_posZmins.value()[i];
+            double zmax = m_posZmaxs.value()[i];
+            z = CLHEP::RandFlat::shoot(zmin, zmax);
+        }
+
+        if (i<m_posRmins.value().size() and i<m_posRmaxs.value().size()) {
+            double rmin = fabs(m_posRmins.value()[i]);
+            double rmax = fabs(m_posRmaxs.value()[i]);
+
+            while (true) {
+                double x_ = CLHEP::RandFlat::shoot(-rmax, rmax);
+                double y_ = CLHEP::RandFlat::shoot(-rmax, rmax);
+                double r_ = std::sqrt(x_*x_+y_*y_);
+
+                if (rmin <= r_ && r_ <= rmax) {
+                    x = x_;
+                    y = y_;
+                    break;
+                }
+            }
+        }
+
+
         mcp.setVertex(edm4hep::Vector3d(x,y,z)); 
         // mcp.setEndpoint();
 
diff --git a/Generator/src/GtGunTool.h b/Generator/src/GtGunTool.h
index fa4f8870..d4b8d9f8 100644
--- a/Generator/src/GtGunTool.h
+++ b/Generator/src/GtGunTool.h
@@ -34,14 +34,23 @@ private:
 
     Gaudi::Property<std::vector<std::string>> m_particles{this, "Particles"};
 
+    // For fixed positions
     Gaudi::Property<std::vector<double>> m_positionXs{this, "PositionXs"};
     Gaudi::Property<std::vector<double>> m_positionYs{this, "PositionYs"};
     Gaudi::Property<std::vector<double>> m_positionZs{this, "PositionZs"};
 
+    // For positions
+    Gaudi::Property<std::vector<double>> m_posZmins{this, "PosZMins"};
+    Gaudi::Property<std::vector<double>> m_posZmaxs{this, "PosZMaxs"};
 
+    Gaudi::Property<std::vector<double>> m_posRmins{this, "PosRMins"};
+    Gaudi::Property<std::vector<double>> m_posRmaxs{this, "PosRMaxs"};
+
+    // For energies
     Gaudi::Property<std::vector<double>> m_energymins{this, "EnergyMins"};
     Gaudi::Property<std::vector<double>> m_energymaxs{this, "EnergyMaxs"};
 
+    // For directions
     Gaudi::Property<std::vector<double>> m_thetamins{this, "ThetaMins"};
     Gaudi::Property<std::vector<double>> m_thetamaxs{this, "ThetaMaxs"};
 
-- 
GitLab