AAFHInterface.cc

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#include <framework/logging/Logger.h>
#include <framework/gearbox/Unit.h>
#include <generators/aafh/AAFHInterface.h>
#include <TRandom3.h>
#include <TDatabasePDG.h>
#include <Math/Vector4D.h>
#include <algorithm>
#include <array>
#include <iomanip>
 
using namespace Belle2;
 
extern "C" {
  double aafhdiag36_rndm_(int*)
  {
    return gRandom->Rndm();
  }
  double aafhdiag36_rnf100_(int*)
  {
    return gRandom->Rndm();
  }
  void aafhdiag36_esft_maxed_(const double* weight)
  {
    B2ERROR("AAFH: Maximum weight to small, increase maxFinalWeight to at least "
            << *weight);
  }
  void aafhdiag36_eswe_maxed_(const double* weight)
  {
    B2ERROR("AAFH: Maximum weight to small, increase maxSubGeneratorWeight to at least "
            << *weight);
  }

  extern void aafhdiag36_init_();
  extern void aafhdiag36_event_(int*);
  extern void aafhdiag36_finish_();

  extern struct {
    double eb; 
    double eswe; 
    double esft; 
    double wap[4]; 
    double wbp[4]; 
    double vap[4]; 
  } input_;

  extern struct {
    double w2min; 
    double w2max; 
  } bound_;

  extern struct {
    int iproc; 
    int irejec; 
    int idump[4]; 
    int info; 
    int itot; 
  } input2_;

  extern struct {
    double face; 
    double facl; 
    double facm; 
    double proc; 
  } factor_;

  extern struct {
    double q1[5]; 
    double q2[5]; 
    double q3[5]; 
    double q4[5]; 
    double q5[5]; 
    double q6[5]; 
  } momenz_;

  extern struct {
    double crosec; 
    double ercros; 
  } output_;

  extern struct {
    double xm;  
    double xmu; 
    double xml; 
    double xm2;  
    double xmu2; 
    double xml2; 
  } masses_;

  extern struct {
    double qcharg; 
    double qchrg2; 
    double qchrg3; 
    double qchrg4; 
  } charge_;

  extern struct {
    double swe[4]; 
    double swek[4]; 
    double mwe[4]; 
    double sum; 
    double sumk; 
    double maxwe; 
    int iwe[4]; 
    int igen; 
  } westat_;

  extern struct {
    double sumft; 
    double sumft2; 
    double ftmax; 
    int ieen; 
  } ftstat_;
};
 
 
void AAFHInterface::setMaxTries(int tries)
{
  input2_.itot = tries;
}
 
void AAFHInterface::setGeneratorWeights(const std::vector<double>& weights)
{
  if (weights.size() > 8) {
    B2WARNING("AAFH: more than 8 generator weights supplied, "
              "ignoring the extra weights");
  } else if (weights.size() > 4 && weights.size() < 8) {
    B2WARNING("AAFH: more than 4 but less than 8 generator weights supplied, "
              "ignoring everything after the first four");
  } else if (weights.size() < 4) {
    B2ERROR("AAFH: not enough generator weights, need exactly four or eight values");
  }
  std::copy_n(weights.begin(), 4, input_.wap);
  if (weights.size() >= 8) {
    std::copy_n(weights.begin() + 4, 4, input_.wbp);
  }
}
 
void AAFHInterface::setMaxWeights(double subgeneratorWeight, double finalWeight)
{
  input_.eswe = subgeneratorWeight;
  input_.esft = finalWeight;
}
 
void AAFHInterface::setSupressionLimits(std::vector<double> limits)
{
  if (limits.size() == 1) {
    B2WARNING("AAFH: Only one suppression limit supplied, using same value for all limits");
    factor_.face = limits[0];
    factor_.facm = limits[0];
    factor_.facl = limits[0];
    factor_.proc = limits[0];
  } else if (limits.size() == 4) {
    factor_.face = limits[0];
    factor_.facm = limits[1];
    factor_.facl = limits[2];
    factor_.proc = limits[3];
  } else {
    B2ERROR("AAFH: Wrong number of suppression limits: supply either one or 4 values");
  }
}
 
int AAFHInterface::getMaxTries() const
{
  return input2_.itot;
}
 
std::vector<double> AAFHInterface::getGeneratorWeights() const
{
  std::vector<double> weights(8);
  std::copy_n(input_.wap, 4, weights.begin());
  std::copy_n(input_.wbp, 4, weights.begin() + 4);
  return weights;
}
 
double AAFHInterface::getMaxSubGeneratorWeight() const
{
  return input_.eswe;
}
 
double AAFHInterface::getMaxFinalWeight() const
{
  return input_.esft;
}
 
std::vector<double> AAFHInterface::getSuppressionLimits() const
{
  return {factor_.face, factor_.facm, factor_.facl, factor_.proc};
}
 
double AAFHInterface::getTotalCrossSection() const
{
  return output_.crosec * Unit::nb;
}
 
double AAFHInterface::getTotalCrossSectionError() const
{
  return output_.ercros * Unit::nb;
}
 
void AAFHInterface::initialize(double beamEnergy, EMode mode, ERejection rejection)
{
  input_.eb = beamEnergy;
  //minimum invariant mass is in units of beam energy and needs to be squared
  bound_.w2min = (m_minMass / beamEnergy) * (m_minMass / beamEnergy);
  //set mode and rejection scheme
  input2_.iproc = mode;
  input2_.irejec = rejection;
 
  //initialize the L
  auto* part = TDatabasePDG::Instance()->GetParticle(m_particle.c_str());
  if (!part) {
    B2ERROR("AAFH: Particle '" << m_particle << "' does not exist");
    return;
  }
  //set the mass in units of beam energy
  masses_.xml = part->Mass() / beamEnergy;
  //set the charge but with opposite sign
  charge_.qcharg = -part->Charge() / 3.0;
  //and remember pdg code
  m_particlePDG = part->PdgCode();
 
  aafhdiag36_init_();
}
 
void AAFHInterface::updateParticle(MCParticleGraph::GraphParticle& p, double* q, int pdg, bool isInitial)
{
  // all values in q are in units of beam energy
  // z-direction inversed since AAFH uses different convention
  ROOT::Math::PxPyPzEVector vec(-q[0]*input_.eb, -q[1]*input_.eb, -q[2]*input_.eb, q[3]*input_.eb);
  p.set4Vector(vec);
 
  // mass is multiplied by charge sign so take the absolute
  p.setMass(fabs(q[4])*input_.eb);
  p.setPDG(pdg);
  p.setStatus(MCParticle::c_PrimaryParticle | MCParticle::c_StableInGenerator);
 
  if (isInitial) {
    p.addStatus(MCParticle::c_Initial);
  }
}
 
void AAFHInterface::generateEvent(MCParticleGraph& mpg)
{
  int status {0};
  aafhdiag36_event_(&status);
  if (status != 0) {
    B2ERROR("Failed to generate event after " << status << " tries, giving up");
    return;
  }
 
  //initial beam particles
  auto& p1 = mpg.addParticle();
  auto& p2 = mpg.addParticle();
 
  //generated particles
  auto& p3 = mpg.addParticle();
  auto& p4 = mpg.addParticle();
  auto& p5 = mpg.addParticle();
  auto& p6 = mpg.addParticle();
 
  //array of generated particle pdg codes depends on generated mode
  std::array<int, 4> pdg {{0}};
  switch (input2_.iproc) {
    case c_MuonParticle:
      pdg = {{m_particlePDG, -m_particlePDG, 13, -13}};
      break;
    case c_MuonMuon:
      pdg = {{13, -13, 13, -13}};
      break;
    case c_ElectonMuon:
      pdg = {{11, -11, 13, -13}};
      break;
    case c_ElectronParticle:
      pdg = {{11, -11, m_particlePDG, -m_particlePDG}};
      break;
    case c_ElectonElectron:
      pdg = {{11, -11, 11, -11}};
      break;
  }
 
  updateParticle(p1, momenz_.q1, 11, true);
  updateParticle(p2, momenz_.q2, -11, true);
  updateParticle(p3, momenz_.q3, pdg[0]);
  updateParticle(p4, momenz_.q4, pdg[1]);
  updateParticle(p5, momenz_.q5, pdg[2]);
  updateParticle(p6, momenz_.q6, pdg[3]);
}
 
void AAFHInterface::finish()
{
  aafhdiag36_finish_();
  // The following code just prints the results. If B2RESULT has been compiled
  // out this is unnecessary and cppcheck complains so we enclose this in an
  // ifdef to only be executed if B2RESULT might actually be shown
#ifndef LOG_NO_B2RESULT
  B2RESULT("AAFH Generation finished.");
  B2RESULT("Final cross section:         "
           << std::setprecision(3) << output_.crosec << "+-"
           << std::setprecision(3) << output_.ercros << " nb");
  B2RESULT("Minimum invariant mass:      "
           << sqrt(bound_.w2min)*input_.eb << " GeV");
  B2RESULT("Maximum invariant mass:      "
           << sqrt(bound_.w2max)*input_.eb << " GeV");
  B2RESULT("Maximum subgenerator weight: " << westat_.maxwe << " ("
           << westat_.mwe[0] << ", "
           << westat_.mwe[1] << ", "
           << westat_.mwe[2] << ", "
           << westat_.mwe[3] << ")");
  B2RESULT("Maximum final weight:        " << ftstat_.ftmax);
  //The fortran source says that it's advisable if all subgenerators have the
  //same probability to account for all peaks in the differental cross section
  //and that it runs with the highest efficiency when the maximum weights in
  //all sub generators are equal. So let's calculate appropriate subgenerator
  //weights by looking how many events were generated per sub generator or the
  //maximum weights
  std::array<double, 4> idealWAP;
  std::array<double, 4> idealWBP;
  //find the first non-zero weight
  int reference = 0;
  for (int i = 0; i < 4; ++i) {
    if (input_.wap[i] != 0) {
      reference = i;
      break;
    }
  }
  //and use that es reference
  const double w0 = westat_.mwe[reference] * input_.wbp[reference];
  const double n0 = westat_.iwe[reference] / input_.wap[reference];
  //now calculate the relative weights we need to get
  //1) same maximum weight which is inversely proportional to to wbp
  //2) same number of generated events which is proportional to wap
  for (int i = 0; i < 4; ++i) {
    if (input_.wap[i] == 0) {
      idealWBP[i] = input_.wbp[i];
      idealWAP[i] = 0;
    } else {
      idealWBP[i] = (westat_.mwe[i] * input_.wbp[i]) / w0;
      idealWAP[i] = n0 * input_.wap[i] / westat_.iwe[i];
      //and normalize to 1.0
      idealWAP[i] /= idealWAP[reference];
    }
  }
 
  B2RESULT("To get the same maximum event weight and chance for each sub "
           "generator use these parameters:");
  B2RESULT("   'maxFinalWeight': " << std::setprecision(3)
           << ftstat_.ftmax * 1.1 << ",");
  B2RESULT("   'maxSubgeneratorWeight': " << std::setprecision(3)
           << westat_.maxwe * 1.1 << ",");
  B2RESULT("   'subgeneratorWeights': [" << std::scientific
           << std::setprecision(3) << idealWAP[0] << ", "
           << std::setprecision(3) << idealWAP[1] << ", "
           << std::setprecision(3) << idealWAP[2] << ", "
           << std::setprecision(3) << idealWAP[3] << ", ");
  B2RESULT("                           " << std::scientific
           << std::setprecision(3) << idealWBP[0] << ", "
           << std::setprecision(3) << idealWBP[1] << ", "
           << std::setprecision(3) << idealWBP[2] << ", "
           << std::setprecision(3) << idealWBP[3] << "],");
#endif
}