Belle II Software  release-05-01-25
ParticleGun.cc
1 /**************************************************************************
2  * BASF2 (Belle Analysis Framework 2) *
3  * Copyright(C) 2010 - Belle II Collaboration *
4  * *
5  * Author: The Belle II Collaboration *
6  * Contributors: Susanne Koblitz, Martin Ritter *
7  * *
8  * This software is provided "as is" without any warranty. *
9  **************************************************************************/
10 
11 #include <framework/logging/Logger.h>
12 #include <framework/gearbox/Unit.h>
13 #include <generators/particlegun/ParticleGun.h>
14 
15 #include <TRandom.h>
16 
17 #include <cmath>
18 #include <limits>
19 
20 using namespace std;
21 using namespace Belle2;
22 
23 namespace {
39  double randomPolyline(size_t n, const double* x, const double* y)
40  {
41  //Calculate the area of the step function which envelops the polyline and
42  //the area in each segment.
43  std::vector<double> weights(n - 1);
44  double sumw(0);
45  for (size_t i = 0; i < n - 1; ++i) {
46  weights[i] = (x[i + 1] - x[i]) * max(y[i], y[i + 1]);
47  sumw += weights[i];
48  }
49  while (true) {
50  //Generate x distributed according to the step function. This can be done
51  //with 100% efficiency
52  double weight = gRandom->Uniform(0, sumw);
53  size_t segment(0);
54  for (; segment < n - 1; ++segment) {
55  weight -= weights[segment];
56  if (weight <= 0) break;
57  }
58  const double x1 = x[segment];
59  const double x2 = x[segment + 1];
60  const double xx = gRandom->Uniform(x1, x2);
61  //Now calculate y(x) using line between (x1,y1) and (x2,y2)
62  const double y1 = y[segment];
63  const double y2 = y[segment + 1];
64  const double yy = (y2 == y1) ? y1 : y1 + (xx - x1) / (x2 - x1) * (y2 - y1);
65  //Generate a random y inside the step function
66  const double randY = gRandom->Uniform(0, max(y1, y2));
67  //And accept the x value if randY lies below the line
68  if (randY < yy) return xx;
69  //Otherwise repeat generation of x and y
70  }
71  }
72 }
73 
74 
75 double ParticleGun::generateValue(EDistribution dist, const vector<double>& params)
76 {
77  double rand(0);
78  switch (dist) {
79  case c_fixedValue:
80  return params[0];
81  case c_uniformDistribution:
82  case c_uniformPtDistribution:
83  return gRandom->Uniform(params[0], params[1]);
84  case c_inversePtDistribution:
85  return 1. / gRandom->Uniform(1. / params[1], 1. / params[0]);
86  case c_uniformCosDistribution:
87  return acos(gRandom->Uniform(cos(params[0]), cos(params[1])));
88  case c_uniformLogDistribution:
89  case c_uniformLogPtDistribution:
90  return exp(gRandom->Uniform(log(params[0]), log(params[1])));
91  case c_normalDistribution:
92  case c_normalPtDistribution:
93  return gRandom->Gaus(params[0], params[1]);
94  case c_normalCosDistribution:
95  return acos(gRandom->Gaus(params[0], params[1]));
96  case c_discreteSpectrum:
97  case c_discretePtSpectrum:
98  //Create weighted discrete distribution
99  //First we sum all the weights (second half of the array)
100  for (size_t i = params.size() / 2; i < params.size(); ++i) rand += params[i];
101  //Then generate a random variable between 0 and sum of weights
102  rand = gRandom->Uniform(0, rand);
103  //Go over the weights again and subtract from generated value
104  for (size_t i = params.size() / 2; i < params.size(); ++i) {
105  rand -= params[i];
106  //If the value is negative we found the correct bin, return that value
107  if (rand <= 0) return params[i - (params.size() / 2)];
108  }
109  B2FATAL("Something wrong with picking fixed spectra values");
110  case c_polylineDistribution:
111  case c_polylinePtDistribution:
112  return randomPolyline(params.size() / 2, params.data(), params.data() + params.size() / 2);
113  case c_polylineCosDistribution:
114  return acos(randomPolyline(params.size() / 2, params.data(), params.data() + params.size() / 2));
115  default:
116  B2FATAL("Unknown distribution");
117  }
118  return 0;
119 }
120 
121 bool ParticleGun::generateEvent(MCParticleGraph& graph)
122 {
123  //generate the event vertex (possible the same for all particles in event)
124  double vx = generateValue(m_params.xVertexDist, m_params.xVertexParams);
125  double vy = generateValue(m_params.yVertexDist, m_params.yVertexParams);
126  double vz = generateValue(m_params.zVertexDist, m_params.zVertexParams);
127 
128  // Time offset
129  double timeOffset = generateValue(m_params.timeDist, m_params.timeParams);
130 
131  //Determine number of tracks
132  int nTracks = static_cast<int>(m_params.nTracks);
133  if (m_params.nTracks <= 0) {
134  nTracks = m_params.pdgCodes.size();
135  } else if (m_params.varyNumberOfTracks) {
136  nTracks = gRandom->Poisson(m_params.nTracks);
137  }
138 
139  //Make list of particles
140  double firstMomentum;
141  for (int i = 0; i < nTracks; ++i) {
143  p.setStatus(MCParticle::c_PrimaryParticle);
144  if (m_params.pdgCodes.size() == 1) {
145  //only one PDGcode available, always take this one
146  p.setPDG(m_params.pdgCodes[0]);
147  } else if (m_params.nTracks <= 0) {
148  //0 or negative nTracks, take the ids sequentially
149  p.setPDG(m_params.pdgCodes[i]);
150  } else {
151  //else choose randomly one of the available codes
152  int index = static_cast<int>(gRandom->Uniform(m_params.pdgCodes.size()));
153  p.setPDG(m_params.pdgCodes[index]);
154  }
155  p.setMassFromPDG();
156  p.setFirstDaughter(0);
157  p.setLastDaughter(0);
158 
159  //lets generate the momentum vector:
160  double momentum = generateValue(m_params.momentumDist, m_params.momentumParams);
161  // if the fixed momentum option enabled, either store the generated momentum for the first
162  // particle or reuse for all others
163  if (m_params.fixedMomentumPerEvent) {
164  i == 0 ? firstMomentum = momentum : momentum = firstMomentum;
165  }
166  double phi = generateValue(m_params.phiDist, m_params.phiParams);
167  double theta = generateValue(m_params.thetaDist, m_params.thetaParams);
168 
169  double pt = momentum * sin(theta);
170  if (m_params.momentumDist == c_uniformPtDistribution || m_params.momentumDist == c_normalPtDistribution ||
171  m_params.momentumDist == c_inversePtDistribution || m_params.momentumDist == c_polylinePtDistribution ||
172  m_params.momentumDist == c_uniformLogPtDistribution || m_params.momentumDist == c_discretePtSpectrum) {
173  //this means we are actually generating the Pt and not the P, so exchange values
174  pt = momentum;
175  momentum = (sin(theta) > 0) ? (pt / sin(theta)) : numeric_limits<double>::max();
176  }
177 
178  double pz = momentum * cos(theta);
179  double px = pt * cos(phi);
180  double py = pt * sin(phi);
181  double m = p.getMass();
182  double e = sqrt(momentum * momentum + m * m);
183 
184  p.setMomentum(px, py, pz);
185  p.setEnergy(e);
186  p.setProductionVertex(vx, vy, vz);
187  p.addStatus(MCParticle::c_StableInGenerator);
188  //Particle is stable in generator. We could use MCParticleGraph options to
189  //do this automatically but setting it here makes the particle correct
190  //independent of the options
191  p.setDecayTime(numeric_limits<double>::infinity());
192 
193  // set time offset to check fit bias in e.g. the ECL waveform fits
194  p.setProductionTime(timeOffset);
195 
196  if (m_params.independentVertices) {
197  //If we have independent vertices, generate new vertex for next particle
198  vx = generateValue(m_params.xVertexDist, m_params.xVertexParams);
199  vy = generateValue(m_params.yVertexDist, m_params.yVertexParams);
200  vz = generateValue(m_params.zVertexDist, m_params.zVertexParams);
201  }
202  }// end loop over particles in event
203 
204  return true;
205 }
206 
207 bool ParticleGun::setParameters(const Parameters& p)
208 {
209  //checking that all parameters are useful is the most complex part of the
210  //whole particle gun. We need to disallow distributions for certain variables
211  //and check the number of parameters and their validity
212 
213  //As long as this is true, no error has been found. We could just return but
214  //it's so much nicer to display all possible errors at once so the user does
215  //not need to iterate over this all the time
216  bool ok(true);
217 
218  //Make an enum -> name mapping for nice error messages
219  std::map<EDistribution, std::string> distributionNames = {
220  {c_fixedValue, "fixed"},
221  {c_uniformDistribution, "uniform"},
222  {c_uniformPtDistribution, "uniformPt"},
223  {c_uniformCosDistribution, "uniformCos"},
224  {c_uniformLogDistribution, "uniformLog"},
225  {c_uniformLogPtDistribution, "uniformLogPt"},
226  {c_normalDistribution, "normal"},
227  {c_normalPtDistribution, "normalPt"},
228  {c_normalCosDistribution, "normalCos"},
229  {c_inversePtDistribution, "inversePt"},
230  {c_polylineDistribution, "polyline"},
231  {c_polylinePtDistribution, "polylinePt"},
232  {c_polylineCosDistribution, "polylineCos"},
233  {c_discreteSpectrum, "discrete"},
234  {c_discretePtSpectrum, "discretePt"}
235  };
236  //Small helper lambda to get the distribution by name
237  auto getDist = [&p](const std::string & dist) -> EDistribution {
238  if (dist == "momentum") return p.momentumDist;
239  if (dist == "xVertex") return p.xVertexDist;
240  if (dist == "yVertex") return p.yVertexDist;
241  if (dist == "zVertex") return p.zVertexDist;
242  if (dist == "theta") return p.thetaDist;
243  if (dist == "phi") return p.phiDist;
244  if (dist == "timeOffset") return p.timeDist;
245  throw std::runtime_error("wrong parameter");
246  };
247  //Small helper lambda to get the parameters by name
248  auto getPars = [&p](const std::string & dist) -> const std::vector<double>& {
249  if (dist == "momentum") return p.momentumParams;
250  if (dist == "xVertex") return p.xVertexParams;
251  if (dist == "yVertex") return p.yVertexParams;
252  if (dist == "zVertex") return p.zVertexParams;
253  if (dist == "theta") return p.thetaParams;
254  if (dist == "phi") return p.phiParams;
255  if (dist == "timeOffset") return p.timeParams;
256  throw std::runtime_error("wrong parameter");
257  };
258  //Small helper lambda to produce a nice error message and set the error flag
259  //if the distribution is excluded.
260  auto excludeDistribution = [&](const std::string & dist, EDistribution excluded) {
261  if (getDist(dist) == excluded) {
262  B2ERROR(distributionNames[excluded] << " is not allowed for " << dist << " generation");
263  ok = false;
264  }
265  };
266 
267  //Exclude some distributions
268  for (auto dist : {"momentum", "xVertex", "yVertex", "zVertex"}) {
269  excludeDistribution(dist, c_uniformCosDistribution);
270  excludeDistribution(dist, c_normalCosDistribution);
271  excludeDistribution(dist, c_polylineCosDistribution);
272  }
273  for (auto dist : {"xVertex", "yVertex", "zVertex", "phi", "theta", "timeOffset"}) {
274  excludeDistribution(dist, c_uniformPtDistribution);
275  excludeDistribution(dist, c_uniformLogPtDistribution);
276  excludeDistribution(dist, c_normalPtDistribution);
277  excludeDistribution(dist, c_inversePtDistribution);
278  excludeDistribution(dist, c_polylinePtDistribution);
279  }
280 
281  //Check that we have some particle ids to generate
282  if (p.pdgCodes.empty()) {
283  B2ERROR("No pdgCodes specified for generation");
284  ok = false;
285  }
286 
287  //Check minimum numbers of parameters
288  for (auto par : {"momentum", "xVertex", "yVertex", "zVertex", "theta", "phi", "timeOffset"}) {
289  const EDistribution dist = getDist(par);
290  size_t minParams = (dist == c_fixedValue) ? 1 : 2;
291  if (dist == c_polylineDistribution || dist == c_polylinePtDistribution || dist == c_polylineCosDistribution) minParams = 4;
292  const size_t hasParams = getPars(par).size();
293  if (hasParams < minParams) {
294  B2ERROR(par << " generation: " << distributionNames[dist]
295  << " distribution requires at least " << minParams
296  << " parameters, " << hasParams << " given");
297  ok = false;
298  }
299  //Check even number of parameters for discrete or polyline distributions
300  if (dist == c_discreteSpectrum || dist == c_discretePtSpectrum ||
301  dist == c_polylineDistribution || dist == c_polylinePtDistribution || dist == c_polylineCosDistribution) {
302  if ((hasParams % 2) != 0) {
303  B2ERROR(par << " generation: " << distributionNames[dist] << " requires an even number of parameters");
304  ok = false;
305  } else if (dist == c_discreteSpectrum || dist == c_discretePtSpectrum || hasParams >= 4) {
306  //Check wellformedness of polyline pdf: ascending x coordinates and positive y coordinates with at least one nonzero value
307  //Discrete spectrum only requires positive weights, no sorting needed
308  const std::vector<double>& pvec = getPars(par);
309  const std::string parname = (dist == c_discreteSpectrum || dist == c_discretePtSpectrum)
310  ? "weight" : "y coordinate";
311  //Check for sorting for polylines
312  if ((dist != c_discreteSpectrum) && (dist != c_discretePtSpectrum)) {
313  for (size_t i = 0; i < (hasParams / 2) - 1; ++i) {
314  if (pvec[i] > pvec[i + 1]) {
315  B2ERROR(par << " generation: " << distributionNames[dist] << " requires x coordinates in ascending order");
316  ok = false;
317  }
318  }
319  }
320  //Check that values are non-negative and at least one value is positive
321  auto minmax = std::minmax_element(pvec.begin() + hasParams / 2, pvec.end());
322  if (*minmax.first < 0) {
323  B2ERROR(par << " generation: " << distributionNames[dist] << " requires "
324  << parname << "s to be non-negative");
325  ok = false;
326  }
327  if (*minmax.second <= 0) {
328  B2ERROR(par << " generation: " << distributionNames[dist] << " requires at least one "
329  << parname << " to be positive");
330  ok = false;
331  }
332  }
333  }
334  }
335 
336  //Finally check that we do not have any problems with the inverse
337  if (p.momentumDist == c_inversePtDistribution && p.momentumParams.size() >= 2) {
338  if (p.momentumParams[0] == 0 || p.momentumParams[1] == 0) {
339  B2ERROR("inversePt distribution does not allow zero momentum");
340  ok = false;
341  }
342  }
343 
344  //If everything is ok, set the new parameters, else return false
345  if (ok) { m_params = p; }
346  return ok;
347 }
Belle2::MCParticleGraph
Class to build, validate and sort a particle decay chain.
Definition: MCParticleGraph.h:48
Belle2::ParticleGun::Parameters
Struct to keep all necessary parameters for the particle gun.
Definition: ParticleGun.h:75
Belle2
Abstract base class for different kinds of events.
Definition: MillepedeAlgorithm.h:19
Belle2::MCParticleGraph::addParticle
GraphParticle & addParticle()
Add new particle to the graph.
Definition: MCParticleGraph.h:305
Belle2::ParticleGun::EDistribution
EDistribution
enum containing all known distributions available for generation of values
Definition: ParticleGun.h:39
Belle2::MCParticleGraph::GraphParticle
Class to represent Particle data in graph.
Definition: MCParticleGraph.h:86