CRTree.H

Go to the documentation of this file.

 
#ifndef _CRTREE_H_
#define _CRTREE_H_
 
#include <cassert>
#include <vector>
#include "CRPacket.H"
#include "CRSource.H"
#include "CRTreeNode.H"
#include "FieldQty.H"
#include "Geometry.H"
#include "../Gas/Gas.H"
#include "../Losses/Ionization.H"
#include "../Losses/Losses.H"
#include "../MPI/CRTreeCommunicator.H"
#include "../MPI/MPIUtil.H"
#include "../Propagation/Propagation.H"
#include "../Utils/RealBox.H"
#include "../Utils/RealTensor2.H"
#include "../Utils/RngThread.H"
#include "../Utils/Types.H"
#include "../Utils/Vec3.H"
 
namespace criptic {
  
  class CRTree {
 
  public:
    
    // Constructor
 
    CRTree(const ParmParser& pp,
       const Geometry& geom_,
       const gas::Gas& gasBG_,
       const propagation::Propagation& prop_,
       const Losses& loss_,
       RngThread& rng_);
 
    // Methods to add or remove packets and sources to the tree
 
    void addPacketsAndSources(const std::vector<RealVec>& x_,
                  const std::vector<CRPacket>& pd_,
                  const std::vector<RealVec>& xSrc_,
                  const std::vector<CRSource>& src_);
 
    
    void addPackets(const std::vector<RealVec>& x_,
            const std::vector<CRPacket>& pd_);
 
    void addSources(const std::vector<RealVec>& xSrc_,
            const std::vector<CRSource>& src_);
 
    void deletePacketsAndSources(std::vector<IdxType>& pIdx,
                 std::vector<IdxType>& sIdx,
                 const Real t);
 
    void deletePackets(std::vector<IdxType>& pIdx, const Real t) {
      std::vector<IdxType> sIdx;
      deletePacketsAndSources(pIdx, sIdx, t);
    }
 
    void deleteSources(std::vector<IdxType>& sIdx) {
      std::vector<IdxType> pIdx;
      Real t = 0;
      deletePacketsAndSources(pIdx, sIdx, t);
    }    
    
    // Methods to retrieve data from the tree
 
    // Counts of objects
    IdxType nPacket() const { return totPacket; }
    IdxType nPacketLoc() const { return x.size() + xBuf.size(); }
    IdxType nDeletedPacket() const { return totDeletedPacket; }
    IdxType nSource() const { return totSrc; }
    IdxType nSourceLoc() const { return src.size(); }
    IdxType getUniqueID() const { return uniqueID; }
    void setUniqueID(const IdxType uniqueID_) { uniqueID = uniqueID_; }
 
    // Packets
    std::vector<RealVec>& getPacketPos() { return x; }
    const std::vector<RealVec>& getPacketPos() const { return x; }
    RealVec& getPacketPos(const IdxType i) {
      assert(i < nPacketLoc());
      if (i < x.size()) return x[i];
      else return xBuf[i-x.size()];
    }
    const RealVec& getPacketPos(const IdxType i) const {
      assert(i < nPacketLoc());
      if (i < x.size()) return x[i];
      else return xBuf[i-x.size()];
    }
    std::vector<CRPacket>& getPacketData() { return pd; }
    const std::vector<CRPacket>& getPacketData() const { return pd; }
    CRPacket& getPacketData(const IdxType i) {
      assert(i < pd.size() + pdBuf.size());
      if (i < pd.size()) return pd[i];
      else return pdBuf[i - pd.size()];
    }
    const CRPacket& getPacketData(const IdxType i) const {
      assert(i < pd.size() + pdBuf.size());
      if (i < pd.size()) return pd[i];
      else return pdBuf[i - pd.size()];
    }
 
    // Deleted packets
    std::vector<RealVec>& getDeletedPacketPos() { return xDel; }
    const std::vector<RealVec>& getDeletedPacketPos() const { return xDel; }
    RealVec& getDeletedPacketPos(const IdxType i) {
      assert(i < xDel.size());
      return xDel[i];
    }
    const RealVec& getDeletedPacketPos(const IdxType i) const {
      assert(i < xDel.size());
      return xDel[i];
    }
    std::vector<CRPacket>& getDeletedPacketData() { return pdDel; }
    const std::vector<CRPacket>& getDeletedPacketData() const { return pdDel; }
    CRPacket& getDeletedPacketData(const IdxType i) {
      assert(i < pdDel.size());
      return pdDel[i];
    }
    const CRPacket& getDeletedPacketData(const IdxType i) const {
      assert(i < pdDel.size());
      return pdDel[i];
    }
    std::vector<Real>& getDeletedPacketTime() { return tDel; }
    const std::vector<Real>& getDeletedPacketTime() const { return tDel; }
    Real& getDeletedPacketTime(const IdxType i) {
      return tDel[i];
    }
    const Real& getDeletedPacketTime(const IdxType i) const {
      return tDel[i];
    }
 
    // Sources
    std::vector<RealVec>& getSourcePos() { return xSrc; }
    const std::vector<RealVec>& getSourcePos() const { return xSrc; }
    RealVec& getSourcePos(const IdxType i) {
      assert(i < xSrc.size());
      return xSrc[i];
    }
    const RealVec& getSourcePos(const IdxType i) const {
      assert(i < xSrc.size());
      return xSrc[i];
    }
    std::vector<CRSource>& getSourceData() { return src; }
    const std::vector<CRSource>& getSourceData() const { return src; }
    CRSource& getSourceData(const IdxType i) {
      assert(i < src.size());
      return src[i];
    }
    const CRSource& getSourceData(const IdxType i) const {
      assert(i < src.size());
      return src[i];
    }
 
    const std::vector<CRSource>& getSourceBufData() const { return srcBuf; }
    
 
    // Scale lengths
    Real& getPacketScale(const IdxType i) {
      assert(i < nPacketLoc());
      if (!needFQ) return h[0];
      else if (i < h.size()) return h[i];
      else return hBuf[i-h.size()];
    }
 
    // Field quantities
    std::vector<FieldQty>& getPacketQty() { return qty; }
    const std::vector<FieldQty>& getPacketQty() const { return qty; }
    FieldQty& getPacketQty(const IdxType i) {
      assert(i < nPacketLoc());
      if (!needFQ) return qty[0];
      else if (i < qty.size()) return qty[i];
      else return qtyBuf[i-qty.size()];
    }
    const FieldQty& getPacketQty(const IdxType i) const {
      assert(i < nPacketLoc());
      if (!needFQ) return qty[0];
      else if (i < qty.size()) return qty[i];
      else return qtyBuf[i-qty.size()];
    }
 
    // Field quantity gradients
    std::vector<FieldQtyGrad>& getPacketQtyGrad() { return qtyGrad; }
    const std::vector<FieldQtyGrad>& getPacketQtyGrad() const {
      return qtyGrad;
    }
    FieldQtyGrad& getPacketQtyGrad(const IdxType i) {
      assert(i < nPacketLoc());
      if (!needFQG) return qtyGrad[0];
      else if (i < qtyGrad.size()) return qtyGrad[i];
      else return qtyGradBuf[i-qtyGrad.size()];
    }
    const FieldQtyGrad& getPacketQtyGrad(const IdxType i) const {
      assert(i < nPacketLoc());
      if (!needFQG) return qtyGrad[0];
      else if (i < qtyGrad.size()) return qtyGrad[i];
      else return qtyGradBuf[i-qtyGrad.size()];
    }
 
    // Loss information
 
    const Losses& getLoss() const { return loss; }
    
    void getLossRates(const Real t,
              std::vector<Real>& dpdt,
              std::vector<Real>& lossRate) const {
      dpdt.resize(pd.size());
      lossRate.resize(pd.size());
#ifdef _OPENMP
#pragma omp parallel for
#endif
      for (IdxType i = 0; i < pd.size(); i++) {
    Real wSec;
    Losses::MechMatrix<Real> pSec;
    gas::GasData gd = gasBG.gasData(x[i], t);
    loss.computeLoss(gd, pd[i], dpdt[i], lossRate[i], wSec, pSec);
      }
    }
 
    void getLossRates(const IdxType i,
              const Real t,
              Real& dpdt,
              Real& lossRate) const {
      const RealVec& x_ = getPacketPos(i);
      const CRPacket& pd_ = getPacketData(i);
      gas::GasData gd = gasBG.gasData(x_, t);
      Losses::MechMatrix<Real> pSec;
      Real wSec;
      loss.computeLoss(gd, pd_, dpdt, lossRate, wSec, pSec);
    }
    
    void getLossRatesMech(const Real t,
              std::vector<Losses::MechArr<Real> >& dpdt,
              std::vector<Losses::MechArr<Real> >& lossRate)
      const {
      dpdt.resize(pd.size());
      lossRate.resize(pd.size());
#ifdef _OPENMP
#pragma omp parallel for
#endif
      for (IdxType i = 0; i < pd.size(); i++) {
    gas::GasData gd = gasBG.gasData(x[i], t);
    dpdt[i] = loss.dpdtMech(gd, pd[i]);
    lossRate[i] = loss.lossRateMech(gd, pd[i]);
      }
    }
 
    void getLossRatesMech(const IdxType i,
              const Real t,
              Losses::MechArr<Real>& dpdt,
              Losses::MechArr<Real>& lossRate) const {
      const RealVec& x_ = getPacketPos(i);
      const CRPacket& pd_ = getPacketData(i);
      gas::GasData gd = gasBG.gasData(x_, t);
      dpdt = loss.dpdtMech(gd, pd_);
      lossRate = loss.lossRateMech(gd, pd_);
    }
 
    // Radiation from packets
 
    std::vector<Real> photLum(const Real t,
                  const std::vector<Real>& en) const {
      IdxType packetSize = lossMech::nMech * en.size();
      std::vector<Real> out(pd.size() * packetSize);
#ifdef _OPENMP
#pragma omp parallel for
#endif
      for (IdxType i = 0; i < pd.size(); i++) {
    gas::GasData gd = gasBG.gasData(x[i], t);
    std::vector<Real> dLdE = loss.photLum(gd, pd[i], en);
    memcpy(out.data() + i*packetSize,
           dLdE.data(),
           packetSize * sizeof(Real));
      }
      return out;
    }
 
    std::vector<Real> photLumSummary(const Real t,
                     const std::vector<Real>& en,
                     const int rank = 0) const;
 
    // Ionization from packets
    
    std::vector<Real> ionRate(const Real t) const {
      std::vector<Real> out(pd.size() * ionData::nTarget);
#ifdef _OPENMP
#pragma omp parallel for
#endif
      for (IdxType i = 0; i < pd.size(); i++) {
    gas::GasData gd = gasBG.gasData(x[i], t);
    out[i * ionData::nTarget] =
      gd.xH0 * gd.nH() * loss.ionRate(pd[i], ionData::targetH);
    out[i * ionData::nTarget + 1] =
      gd.xH2() * gd.nH() * loss.ionRate(pd[i], ionData::targetH2);
    out[i * ionData::nTarget + 2] =
      gd.xHe0 * gd.nH() * loss.ionRate(pd[i], ionData::targetHe);
      }
      return out;
    }
 
#ifndef TRACK_PITCH_ANGLE
    std::vector<Real> computePdot(const Real t) const;
#endif
 
    // Pass-through routines allowing getting and setting of the
    // RNG state
 
    std::string getRngState() const {
      return rng.getState();
    }
 
    void setRngState(std::string &buf) const {
      rng.setState(buf);
    }
    
    // Method to flush deleted packets from memory
    
    void clearDeletedPackets() {
      xDel.resize(0);
      pdDel.resize(0);
    }
 
    // Public constants
 
    static constexpr int
    leafSize = TREE_LEAF_SIZE;  
    static constexpr int
    nRigidBin = TREE_RIGIDITY_BINS; 
    // Methods that form part of the update cycle
 
    void rebuild();
 
    void fillFieldQty() {
      if (x.size() == 0) return; // If no packets, return immediately
      std::vector<IdxType> start(leafList.size());
      std::vector<IdxType> count(leafList.size());
      for (IdxType i = 0; i < leafList.size(); i++) {
    start[i] = nodes[leafList[i]].start;
    count[i] = nodes[leafList[i]].count;
      }
      fillFieldQty(start, count, leafList);
    }
 
    void injectPackets(const Real tStart,
               const Real tStop);
 
    Real advancePackets(const Real tStart,
            const Real tStop);
        
  private:
 
    // Steps carried out by rebuild
 
    void setupDomain();
 
    void buildGlobalTree();
    
    void buildLocalTree();
 
    void computeSplit(GlobalNode& nd);
    
    template <class U>
    IdxType partitionGlobalNode(const GlobalNode& nd,
                const IdxType start,
                const IdxType count,
                std::vector<RealVec>& x_,
                std::vector<U>& dat);
 
    void partitionNode(LocalNode& nd);
 
    void sumTree();
 
#ifdef ENABLE_MPI
    void exchange();
#endif
    
    // Methods used by fillFieldQty
 
    void fillFieldQty(std::vector<IdxType>& start,
              std::vector<IdxType>& count,
              std::vector<IdxType>& leafList_);
    
    void fillFieldQtyRange(const IdxType start,
               const IdxType count,
               const RealVec x_[],
               const Real R[],
               const RealTensor2& h2Inv,
               const RealTensor2& h2InvGrad,
               const Real& hScale,
               FieldQty q[],
               FieldQtyGrad qGrad[]
#ifdef ENABLE_MPI
               , const int nrank = 0,
               const FieldQty* qOwner = nullptr
#endif
               );
 
    void addNodeToField(const IdxType& nodeIdx,
            const IdxType& start,
            const IdxType& count,
            const RealVec x_[],
            const Real R[],
            const RealTensor2& h2Inv,
            const RealTensor2& h2InvGrad,
            FieldQty q[],
            FieldQty qErr[],
            FieldQtyGrad qGrad[],
            FieldQtyGrad qGradErr[]);
 
    void nodeWgtLimits(const RealVec& x_,
               const RealBox& bnd,
               const RealTensor2& h2Inv,
               const RealTensor2& h2InvGrad,
               Real& wMin,
               Real& wMax,
               RealVec& wGradMin,
               RealVec& wGradMax) const;
    
    // Internal data
    
    // References to other objects
    const Geometry& geom;    
    const gas::Gas& gasBG;   
    const propagation::Propagation& prop; 
    const Losses& loss;      
    RngThread& rng;          
    // Helper classes
#ifdef ENABLE_MPI
    CRTreeCommunicator comm; 
    friend class CRTreeCommunicator;
#endif
    
    // Internal counters
    IdxType uniqueID;        
    // Problem run parameters read or computed from input deck
    int verbosity;           
    int globTreeExtLev;      
    int partitionSamples;    
    int nNgbEff;             
    int equalSampling;       
    Real kdeTol;             
    Real hDither;            
    Real qSamp;              
    Real packetRate;         
    Real cMinWgt;            
    Real cStep;              
    Real errTol;             
    Real wFracMin;           
    Real pMin;               
    Real TMin;               
    Real packetDtMin;        
    std::vector<Real> samplingRatio; 
    bool needFQ;             
    bool needFQG;            
    // CR packet data
    IdxType totPacket;             
    std::vector<RealVec> x;        
    std::vector<CRPacket> pd;      
    std::vector<FieldQty> qty;     
    std::vector<Real> h;           
    std::vector<FieldQtyGrad> qtyGrad; 
    std::vector<RealVec> xBuf;     
    std::vector<CRPacket> pdBuf;   
    std::vector<FieldQty> qtyBuf;  
    std::vector<FieldQtyGrad> qtyGradBuf; 
    std::vector<Real> hBuf;        
    // Data on deleted CR packets
    IdxType totDeletedPacket;      
    std::vector<RealVec> xDel;     
    std::vector<CRPacket> pdDel;   
    std::vector<Real> tDel;        
    // CR source data
    IdxType totSrc;                
    Real totSrcWgt;                
    std::vector<Real> totSrcWgtSpec; 
    std::vector<RealVec> xSrc;     
    std::vector<CRSource> src;     
    std::vector<RealVec> xSrcBuf;  
    std::vector<CRSource> srcBuf;  
    std::vector<IdxType> srcLeaf;  
    // Data describing the domain
    RealVec domLo;                 
    RealVec domHi;                 
    Real logR0;                    
    Real logR1;                    
    Real dlogR;                    
    // Data describing the tree
    IdxType globTreeMaxLev;        
    IdxType globTreeNLeaf;         
    IdxType globTreeLeafPtr;       
    IdxType globTreeMyLeafPtr;     
    IdxType globTreeNOwned;        
    std::vector<IdxType> locTreeOffsets; 
    std::vector<IdxType> locTreeMaxLev;  
    std::vector<IdxType> locTreeLeafStart;  
    std::vector<IdxType> locTreeNLeaf;   
    std::vector<IdxType> leafList;       
    std::vector<GlobalNode> globNodes;   
    std::vector<LocalNode> nodes;        
    std::vector<RealTensor2> h2InvLeaf;  
    std::vector<RealTensor2> h2InvGradLeaf;  
    std::vector<Real>  hScaleLeaf; 
  };
 
}
 
#endif
// _CRTREE_H_