MultiPartRain_module.cc

Go to the documentation of this file.

////////////////////////////////////////////////////////////////////////
// Class:       MultiPartRain
// Module Type: producer
// File:        MultiPartRain_module.cc
//
// Generated at Tue Dec 13 15:48:59 2016 by Kazuhiro Terao using artmod
// from cetpkgsupport v1_11_00.
////////////////////////////////////////////////////////////////////////

#include "art/Framework/Core/EDProducer.h"
#include "art/Framework/Core/ModuleMacros.h"
#include "art/Framework/Principal/Event.h"
#include "art/Framework/Principal/Handle.h"
#include "art/Framework/Principal/Run.h"
#include "art/Framework/Principal/SubRun.h"
//#include "art/Utilities/InputTag.h"
#include "fhiclcpp/ParameterSet.h"
#include "messagefacility/MessageLogger/MessageLogger.h"

#include <memory>
#include <vector>
#include <cmath>
#include "cetlib/pow.h"

#include "CLHEP/Random/RandFlat.h"
#include "CLHEP/Random/RandGauss.h"
#include "CLHEP/Random/RandGeneral.h"

#include "TRandom.h"
#include "nurandom/RandomUtils/NuRandomService.h"
#include "larcore/Geometry/Geometry.h"
#include "larcore/CoreUtils/ServiceUtil.h"
#include "larcoreobj/SummaryData/RunData.h"
#include "larcoreobj/SimpleTypesAndConstants/PhysicalConstants.h"

#include "nusimdata/SimulationBase/MCTruth.h"
#include "nusimdata/SimulationBase/MCParticle.h"

#include "TLorentzVector.h"
#include "TDatabasePDG.h"

struct PartGenParam {
    std::vector<int       > pdg;
    std::vector<double    > mass;
    std::array <size_t, 2 > multi;
    std::array <double, 2 > kerange;
    bool use_mom;
    bool direct_inward;
    double weight;
};

class MultiPartRain;

class MultiPartRain : public art::EDProducer {
public:
    explicit MultiPartRain(fhicl::ParameterSet const & p);
    // The destructor generated by the compiler is fine for classes
    // without bare pointers or other resource use.
    ~MultiPartRain();

    // Plugins should not be copied or assigned.
    MultiPartRain(MultiPartRain const &) = delete;
    MultiPartRain(MultiPartRain &&) = delete;
    MultiPartRain & operator = (MultiPartRain const &) = delete;
    MultiPartRain & operator = (MultiPartRain &&) = delete;

    // Required functions.
    void produce(art::Event & e) override;

    void beginRun(art::Run& run) override;

    void GenPosition(double& x, double& y, double& z);

    std::array<double, 3U> extractDirection() const;
    void GenMomentum(const PartGenParam& param, const double& mass, double& px, double& py, double& pz,
            const double x, const double y, const double z);

    std::vector<size_t> GenParticles() const;

private:

    CLHEP::HepRandomEngine& fFlatEngine;
    std::unique_ptr<CLHEP::RandFlat> fFlatRandom;
    std::unique_ptr<CLHEP::RandGauss> fNormalRandom;
    std::unique_ptr<CLHEP::RandGeneral> fCosmicAngleRandom;

    // exception thrower
    void abort(const std::string msg) const;

    // array of particle info for generation
    std::vector<PartGenParam> _param_v;

    // g4 time of generation
    double _t0;
    double _t0_sigma;

    // g4 position
    std::array<double,2> _xrange;
    std::array<double,2> _yrange;
    std::array<double,2> _zrange;

    // Whether to use uniform or cosmic-like angle distribution
    bool _cosmic_distribution;

    // TPC array
    std::vector<std::vector<unsigned short> > _tpc_v;

    // multiplicity constraint
    size_t _multi_min;
    size_t _multi_max;

    // verbosity flag
    unsigned short _debug;
};

void MultiPartRain::abort(const std::string msg) const
{
    std::cerr << "\033[93m" << msg.c_str() << "\033[00m" << std::endl;
    throw std::exception();
}

MultiPartRain::~MultiPartRain()
{ }

MultiPartRain::MultiPartRain(fhicl::ParameterSet const & p)
: EDProducer(p)
, fFlatEngine(art::ServiceHandle<rndm::NuRandomService>()->createEngine(*this, "HepJamesRandom", "GenRain"))
  //, fFlatEngine(art::ServiceHandle<rndm::NuRandomService>()->createEngine(*this, "HepJamesRandom", "Gen", p, "Seed"))
// :
// Initialize member data here.
{

    //
    // Random engine initialization
    //
    fFlatRandom = std::make_unique<CLHEP::RandFlat>(fFlatEngine,0,1);
    fNormalRandom = std::make_unique<CLHEP::RandGauss>(fFlatEngine);

    // x^2 distribution for fCosmicAngleRandom
    // will be used to draw cos(theta)
    const int nbins(100);
    double parent[nbins];
    for (size_t idx = 0; idx < nbins; ++idx) {
        //parent[idx] = cet::square(cos(((float) idx)/nbins * util::pi()));
        parent[idx] = cet::square(((float) idx)/nbins);
    }
    fCosmicAngleRandom = std::make_unique<CLHEP::RandGeneral>(fFlatEngine, parent, nbins);

    produces< std::vector<simb::MCTruth>   >();
    produces< sumdata::RunData, art::InRun >();

    _debug = p.get<unsigned short>("DebugMode",0);

    _t0 = p.get<double>("G4Time");
    _t0_sigma = p.get<double>("G4TimeJitter");
    if(_t0_sigma < 0) this->abort("Cannot have a negative value for G4 time jitter");

    _multi_min = p.get<size_t>("MultiMin");
    _multi_max = p.get<size_t>("MultiMax");
    _cosmic_distribution = p.get<bool>("CosmicDistribution", false);

    _tpc_v = p.get<std::vector<std::vector<unsigned short> > >("TPCRange");
    auto const xrange = p.get<std::vector<double> > ("XRange");
    auto const yrange = p.get<std::vector<double> > ("YRange");
    auto const zrange = p.get<std::vector<double> > ("ZRange");
    auto const part_cfg = p.get<fhicl::ParameterSet>("ParticleParameter");

    _param_v.clear();
    auto const pdg_v      = part_cfg.get<std::vector<std::vector<int>    > > ("PDGCode");
    auto const minmult_v  = part_cfg.get<std::vector<unsigned short> > ("MinMulti");
    auto const maxmult_v  = part_cfg.get<std::vector<unsigned short> > ("MaxMulti");
    auto const weight_v   = part_cfg.get<std::vector<double> > ("ProbWeight");

    auto kerange_v  = part_cfg.get<std::vector<std::vector<double> > > ("KERange");
    auto momrange_v = part_cfg.get<std::vector<std::vector<double> > > ("MomRange");

    if( (kerange_v.empty() && momrange_v.empty()) ||
       (!kerange_v.empty() && !momrange_v.empty()) ) {
        this->abort("Only one of KERange or MomRange must be empty!");
    }

    bool use_mom = false;
    if(kerange_v.empty()){
        kerange_v = momrange_v;
        use_mom = true;
    }
    // sanity check
    if( pdg_v.size() != kerange_v.size() ||
       pdg_v.size() != minmult_v.size() ||
       pdg_v.size() != maxmult_v.size() ||
       pdg_v.size() != weight_v.size() )
        this->abort("configuration parameters have incompatible lengths!");

    // further sanity check (1 more depth for some double-array)
    for(auto const& r : pdg_v    ) { if(              r.empty()  ) this->abort("PDG code not given!");                        }
    for(auto const& r : kerange_v) { if(              r.size()!=2) this->abort("Incompatible legnth @ KE vector!");           }

    for(size_t idx=0; idx<minmult_v.size(); ++idx) {
        if(minmult_v[idx] > maxmult_v[idx]) this->abort("Particle MinMulti > Particle MaxMulti!");
        if(minmult_v[idx] > _multi_max) this->abort("Particle MinMulti > overall MultiMax!");
        if(minmult_v[idx] > _multi_min)
            _multi_min = minmult_v[idx];
    }
    if(_multi_max < _multi_min) this->abort("Overall MultiMax <= overall MultiMin!");

    if(!xrange.empty() && xrange.size() >2) this->abort("Incompatible legnth @ X vector!" );
    if(!yrange.empty() && yrange.size() >2) this->abort("Incompatible legnth @ Y vector!" );
    if(!zrange.empty() && zrange.size() >2) this->abort("Incompatible legnth @ Z vector!" );

    // slight modification from mpv: define the overall volume across specified TPC IDs + range options
    double xmin,xmax,ymin,ymax,zmin,zmax;
    xmin = ymin = zmin =  1.e20;
    xmax = ymax = zmax = -1.e20;
    // Implementation of required member function here.
    auto geop = lar::providerFrom<geo::Geometry>();
    for(auto const& tpc_id : _tpc_v) {
        assert(tpc_id.size() == 2);
        size_t cid = tpc_id[0];
        size_t tid = tpc_id[1];
        auto const& cryostat = geop->Cryostat(cid);
        assert(cryostat.HasTPC(tid));

        auto const& tpc = cryostat.TPC(tid);
        auto const& tpcabox = tpc.ActiveBoundingBox();
        xmin = std::min(tpcabox.MinX(), xmin);
        ymin = std::min(tpcabox.MinY(), ymin);
        zmin = std::min(tpcabox.MinZ(), zmin);
        xmax = std::max(tpcabox.MaxX(), xmax);
        ymax = std::max(tpcabox.MaxY(), ymax);
        zmax = std::max(tpcabox.MaxZ(), zmax);

        if(_debug) {
            std::cout << "Using Cryostat " << tpc_id[0] << " TPC " << tpc_id[1]
                    << " ... X " << xmin << " => " << xmax
                    << " ... Y " << ymin << " => " << ymax
                    << " ... Z " << zmin << " => " << zmax
                    << std::endl;
        }
    }

    // range register
    if(xrange.size()==1) { _xrange[0] = _xrange[1] = xrange[0]; }
    if(yrange.size()==1) { _yrange[0] = _yrange[1] = yrange[0]; }
    if(zrange.size()==1) { _zrange[0] = _zrange[1] = zrange[0]; }
    if(xrange.size()==2) { _xrange[0] = xrange[0]; _xrange[1] = xrange[1]; }
    if(yrange.size()==2) { _yrange[0] = yrange[0]; _yrange[1] = yrange[1]; }
    if(zrange.size()==2) { _zrange[0] = zrange[0]; _zrange[1] = zrange[1]; }

    _xrange[0] = xmin + _xrange[0];
    _xrange[1] = xmax - _xrange[1];
    _yrange[0] = ymin + _yrange[0];
    _yrange[1] = ymax - _yrange[1];
    _zrange[0] = zmin + _zrange[0];
    _zrange[1] = zmax - _zrange[1];

    // check
    assert(_xrange[0] <= _xrange[1]);
    assert(_yrange[0] <= _yrange[1]);
    assert(_zrange[0] <= _zrange[1]);

    if(_debug>0) {
        std::cout<<"Particle generation world boundaries..."<<std::endl
            <<"X " << _xrange[0] << " => " << _xrange[1] << std::endl
            <<"Y " << _yrange[0] << " => " << _yrange[1] << std::endl
            <<"Z " << _zrange[0] << " => " << _zrange[1] << std::endl;
    }


    // register
    //auto db = new TDatabasePDG;
    auto db = TDatabasePDG::Instance();
    bool direct_inward = p.get<bool>("DirectInward", true);
    if (_cosmic_distribution && direct_inward) {
        std::cout << "Cannot satisfy both CosmicDistribution and Inward Directioning at the same time" << std::endl;
        throw std::exception();
    }
    for(size_t idx=0; idx<pdg_v.size(); ++idx) {
        auto const& pdg     = pdg_v[idx];
        auto const& kerange = kerange_v[idx];
        PartGenParam param;
        param.direct_inward = direct_inward;
        param.use_mom    = use_mom;
        param.pdg        = pdg;
        param.kerange[0] = kerange[0];
        param.kerange[1] = kerange[1];
        param.mass.resize(pdg.size());
        param.multi[0]   = minmult_v[idx];
        param.multi[1]   = maxmult_v[idx];
        param.weight     = weight_v[idx];
        for(size_t i=0; i<pdg.size(); ++i)
            param.mass[i] = db->GetParticle(param.pdg[i])->Mass();

        // sanity check
        if(kerange[0]<0 || kerange[1]<0)
            this->abort("You provided negative energy? Fuck off Mr. Trump.");

        // overall range check
        if(param.kerange[0] > param.kerange[1]) this->abort("KE range has no phase space...");

        if(_debug>0) {
            std::cout << "Generating particle (PDG";
            for(auto const& pdg : param.pdg) std::cout << " " << pdg;
            std::cout << ")" << std::endl
            << (param.use_mom ? "    KE range ....... " : "    Mom range ....... ")
            << param.kerange[0] << " => " << param.kerange[1] << " MeV" << std::endl
            << std::endl;
        }

        _param_v.push_back(param);
    }
}

void MultiPartRain::beginRun(art::Run& run)
{
    // grab the geometry object to see what geometry we are using
    art::ServiceHandle<geo::Geometry> geo;

    std::unique_ptr<sumdata::RunData> runData(new sumdata::RunData(geo->DetectorName()));

    run.put(std::move(runData));

    return;
}

std::vector<size_t> MultiPartRain::GenParticles() const {

    std::vector<size_t> result;
    std::vector<size_t> gen_count_v(_param_v.size(),0);
    std::vector<double> weight_v(_param_v.size(),0);
    for(size_t idx=0; idx<_param_v.size(); ++idx)
        weight_v[idx] = _param_v[idx].weight;

    int num_part = (int)(fFlatRandom->fire(_multi_min,_multi_max+1-1.e-10));

    while(num_part) {

        double total_weight = 0;
        for(auto const& v : weight_v) total_weight += v;

        double rval = 0;
        rval = fFlatRandom->fire(0,total_weight);

        size_t idx = 0;
        for(idx = 0; idx < weight_v.size(); ++idx) {
            rval -= weight_v[idx];
            if(rval <=0.) break;
        }

        // register to the output
        result.push_back(idx);

        // if generation count exceeds max, set probability weight to be 0
        gen_count_v[idx] += 1;
        if(gen_count_v[idx] >= _param_v[idx].multi[1])
            weight_v[idx] = 0.;

        --num_part;
    }
    return result;
}

void MultiPartRain::GenPosition(double& x, double& y, double& z) {

  x = fFlatRandom->fire(_xrange[0],_xrange[1]);
  y = fFlatRandom->fire(_yrange[0],_yrange[1]);
  z = fFlatRandom->fire(_zrange[0],_zrange[1]);

  if(_debug>0) {
    std::cout << "Generating a rain particle at ("
        << x << "," << y << "," << z << ")" << std::endl;
  }
}

std::array<double, 3U> MultiPartRain::extractDirection() const {
    double px, py, pz;
    if (_cosmic_distribution) {
        double phi   = fFlatRandom->fire(0, 2 * util::pi());
        // Zenith Angle Theta is in [pi/2, pi]
        double costheta = - 1. * fCosmicAngleRandom->fire();
        double sintheta = sqrt(1 - costheta * costheta);
        pz = cos(phi) * sintheta;
        px = sin(phi) * sintheta;
        py = costheta;
    } else {
        px = fNormalRandom->fire(0, 1);
        py = fNormalRandom->fire(0, 1);
        pz = fNormalRandom->fire(0, 1);
        double p = std::hypot(px, py, pz);
        px = px / p;
        py = py / p;
        pz = pz / p;
    }

    std::array<double, 3U> result = {px, py, pz};
    return result;
}

void MultiPartRain::GenMomentum(const PartGenParam& param, const double& mass, double& px, double& py, double& pz,
    const double x, const double y, const double z) {

    double tot_energy = 0;
    if(param.use_mom)
      tot_energy = sqrt(cet::square(fFlatRandom->fire(param.kerange[0],param.kerange[1])) + cet::square(mass));
    else
      tot_energy = fFlatRandom->fire(param.kerange[0],param.kerange[1]) + mass;

    double mom_mag = sqrt(cet::square(tot_energy) - cet::square(mass));

    /* Generating unit vector with uniform distribution
     * in direction = over the sphere.
     *
     * It is sufficient to draw a normal variable in
     * each direction and normalize.
     *
     * https://mathworld.wolfram.com/SpherePointPicking.html
     */

    if(!param.direct_inward) {
      std::cout<<"No inward directioning..."<<std::endl;
      std::array<double, 3U> p = extractDirection();
      px = p[0]; py = p[1]; pz = p[2];
    }else{
      double sign_x = ( (x - _xrange[0]) < (_xrange[1] - x) ? 1.0 : -1.0 );
      double sign_y = ( (y - _yrange[0]) < (_yrange[1] - y) ? 1.0 : -1.0 );
      double sign_z = ( (z - _zrange[0]) < (_zrange[1] - z) ? 1.0 : -1.0 );

      if(_debug) {
        std::cout<<"Generating XYZ direction sign for a particle at (" << x << "," << y << "," << z << ")" << std::endl
                <<"X: " << _xrange[0] << " => " << _xrange[1] << " ... " << sign_x << std::endl
                <<"Y: " << _yrange[0] << " => " << _yrange[1] << " ... " << sign_y << std::endl
                <<"Z: " << _zrange[0] << " => " << _zrange[1] << " ... " << sign_z << std::endl
                << std::endl;
      }

      while(1) {
        std::array<double, 3U> p = extractDirection();
        px = p[0]; py = p[1]; pz = p[2];
        if( (px * sign_x) >= 0. &&
                (py * sign_y) >= 0. &&
                (pz * sign_z) >= 0. )
            break;
      }

    }
    //std::cout<<"LOGME,"<<phi<<","<<theta<<","<<px<<","<<py<<","<<pz<<std::endl;

    if(_debug>1)
        std::cout << "    Direction : (" << px << "," << py << "," << pz << ")" << std::endl
        << "    Momentum  : " << mom_mag << " [MeV/c]" << std::endl
        << "    Energy    : " << tot_energy << " [MeV/c^2]" << std::endl;
    px *= mom_mag;
    py *= mom_mag;
    pz *= mom_mag;

}

void MultiPartRain::produce(art::Event & e)
{
    if(_debug>0) std::cout << "Processing a new event..." << std::endl;

    std::unique_ptr< std::vector<simb::MCTruth> > mctArray(new std::vector<simb::MCTruth>);

    simb::MCTruth mct;

    mct.SetOrigin(simb::kCosmicRay);

    std::vector<simb::MCParticle> part_v;

    auto const param_idx_v = GenParticles();
    if(_debug)
        std::cout << "Generating" << param_idx_v.size() << " particles..." << std::endl;

    for(size_t idx=0; idx<param_idx_v.size(); ++idx) {
        auto const& param = _param_v[param_idx_v[idx]];
        double px,py,pz;
        // decide which particle
        size_t pdg_index = (size_t)(fFlatRandom->fire(0,param.pdg.size()-1.e-10));
        auto const& pdg  = param.pdg[pdg_index];
        auto const& mass = param.mass[pdg_index];
        if(_debug) std::cout << "  " << idx << "th instance PDG " << pdg << std::endl;
        double x, y, z;
        GenPosition(x,y,z);
        double g4_time = fFlatRandom->fire(_t0 - _t0_sigma/2., _t0 + _t0_sigma/2.);
        TLorentzVector pos(x,y,z,g4_time);
        std::cout<<"Generating momentum..."<<std::endl;
        GenMomentum(param,mass,px,py,pz,x,y,z);
        std::cout<<"done"<<std::endl;
        TLorentzVector mom(px,py,pz,sqrt(cet::square(px)+cet::square(py)+cet::square(pz)+cet::square(mass)));
        simb::MCParticle part(part_v.size(), pdg, "primary", 0, mass, 1);
        part.AddTrajectoryPoint(pos,mom);
        part_v.emplace_back(std::move(part));
    }

    if(_debug) std::cout << "Total number particles: " << mct.NParticles() << std::endl;
    /*
     simb::MCParticle nu(mct.NParticles(), 16, "primary", mct.NParticles(), 0, 0);
     double px=0;
     double py=0;
     double pz=0;
     double en=0;
     for(auto const& part : part_v) {
     px = part.Momentum().Px();
     py = part.Momentum().Py();
     pz = part.Momentum().Pz();
     en = part.Momentum().E();
     }
     TLorentzVector mom(px,py,pz,en);
     nu.AddTrajectoryPoint(pos,mom);

     mct.Add(nu);
     */
    for(auto& part : part_v)
        mct.Add(part);

    mctArray->push_back(mct);

    e.put(std::move(mctArray));
}

DEFINE_ART_MODULE(MultiPartRain)