PartBunch.h

Go to the documentation of this file.

 
#ifndef PARTBUNCH_H
#define PARTBUNCH_H
 
#include <memory>
#include <vector>
 
#include "Algorithms/CoordinateSystemTrafo.h"
#include "Algorithms/Matrix.h"
#include "Algorithms/PartData.h"
#include "Attributes/Attributes.h"
#include "BCHandler.hpp"
#include "Manager/BaseManager.h"
#include "Manager/PicManager.h"
#include "PartBunch/Binning/AdaptBins.h"
#include "PartBunch/BunchStateHandler.h"
#include "PartBunch/FieldContainer.hpp"
#include "PartBunch/FieldSolver.hpp"
#include "PartBunch/LoadBalancer.hpp"
#include "PartBunch/ParticleContainer.hpp"
#include "Physics/Physics.h"
#include "Random/Distribution.h"
#include "Random/InverseTransformSampling.h"
#include "Random/NormalDistribution.h"
#include "Random/Randn.h"
#include "Structure/FieldSolverCmd.h"
#include "Utilities/OpalException.h"
 
class DataSink;  
class Beam;
 
extern Inform* gmsg;
 
using view_type = typename ippl::detail::ViewType<ippl::Vector<double, 3>, 1>::view_type;
 
template <typename T, unsigned Dim>
class BinnedFieldSolver;
 
template <typename T, unsigned Dim>
class PartBunch
    : public ippl::PicManager<
              T, Dim, ParticleContainer<T, Dim>, FieldContainer<T, Dim>, LoadBalancer<T, Dim>> {
public:
    using ParticleContainer_t = ParticleContainer<T, Dim>;
    using FieldContainer_t    = FieldContainer<T, Dim>;
    using BinnedFieldSolver_t = BinnedFieldSolver<T, Dim>;
    using LoadBalancer_t      = LoadBalancer<T, Dim>;
    using Base                = ippl::ParticleBase<
                           ippl::ParticleSpatialLayout<T, Dim>, Kokkos::DefaultExecutionSpace::memory_space>;
 
    using CoordinateSelector_t = typename ParticleBinning::CoordinateSelector<ParticleContainer_t>;
    using GammaSelector_t      = typename ParticleBinning::GammaSelector<ParticleContainer_t>;
    using AdaptBins_t          = typename ParticleBinning::AdaptBinsBase<ParticleContainer_t>;
    using binIndex_t           = typename ParticleContainer_t::bin_index_type;
 
    using BCHandler_t = BCHandler<Dim>;
 
public:
    // --- Shared state (all containers / mesh) ---
 
    double dt_m;                       
    int it_m;                          
    std::string integration_method_m;  
    std::string solver_m;              
    Vector_t<int, Dim> nr_m;           
    int nrZBase_m = 0;  
    Vector_t<double, Dim> origin_m;  
    Vector_t<double, Dim> rmin_m;  
    Vector_t<double, Dim>
            rmax_m;              
    Vector_t<double, Dim> hr_m;  
 
    double lbt_m;                    
    bool isFirstRepartition_m;       
    ippl::NDIndex<Dim> domain_m;     
    std::array<bool, Dim> decomp_m;  
 
private:
    std::vector<bool> pcActive_m;   
    std::vector<bool> pcAtZStop_m;  
    std::vector<std::string> particleNames_m;  
 
    std::vector<double> qi_m;  
    std::vector<double> mi_m;  
 
    const PartData* reference_m = nullptr;  
 
    std::shared_ptr<BunchStateHandler>
            bunchState_m;  
 
    std::shared_ptr<BCHandler_t> bcHandler_m;  
    std::shared_ptr<AdaptBins_t> bins_m;       
    FieldSolverCmd* OPALFieldSolver_m;         
    DataSink* dataSink_m;                      
 
    double t_m;  
 
    std::shared_ptr<VField_t<T, Dim>> Etmp_m;
    std::shared_ptr<VField_t<T, Dim>> Btmp_m;
 
    long long globalTrackStep_m;  
 
    std::unique_ptr<size_t[]>
            globalPartPerNode_m;  
 
    double rmsDensity_m;  
 
public:
    PartBunch(
            std::vector<double> qi, std::vector<double> mi, const std::vector<Beam*>& beams,
            std::vector<size_t> totalParticlesPerBeam, double lbt, std::string integration_method,
            FieldSolverCmd* OPALFieldSolver, DataSink* dataSink);
 
    void bunchUpdate();
 
    void computeBoundsForFieldSolve(Vector_t<double, Dim>& lower, Vector_t<double, Dim>& upper);
 
    void applyGridUpdate(const Vector_t<double, Dim>& lower, const Vector_t<double, Dim>& upper);
 
    void reinitializeGridZ(int nrZ);
 
    void setImageChargeConfiguration(bool enabled, double zPlane);
 
    void setShiftedGreensConfiguration(bool enabled, double zPlane);
 
    void setZeroFacePlaneDumpFrequency(int frequency);
 
    void setZerofaceMaxSteps(int maxSteps);
 
    size_t getTotalNumAllContainers() const {
        size_t total = 0;
        for (const auto& pc : this->getParticleContainers()) {
            if (pc) {
                total += pc->getTotalNum();
            }
        }
        return total;
    }
 
    void setSolver();
 
    void setBins();
 
    void pre_run() override;
 
    void performBunchSanityChecks() const;
 
    void resetPcActive();
 
    bool isPcActive(size_t i) const { return i < pcActive_m.size() && pcActive_m[i]; }
 
    void setPcActive(size_t i) {
        if (i < pcActive_m.size()) {
            pcActive_m[i] = true;
        }
    }
 
    bool pcAtZStop(size_t i) const { return i < pcAtZStop_m.size() && pcAtZStop_m[i]; }
 
    void setPcAtZStop(size_t i);
 
    void refreshPcActiveAfterEmit();
 
    bool anyPcActive() const {
        for (bool a : pcActive_m) {
            if (a) {
                return true;
            }
        }
        return false;
    }
 
    void advance() override {
        throw OpalException(
                "PartBunch::advance",
                "Not used: just exists because ippl::PicManager wants it that way.");
    }
 
    void par2grid() override {
        throw OpalException(
                "PartBunch::par2grid",
                "Not used: just exists because ippl::PicManager wants it that way.");
    }
 
    void grid2par() override {
        throw OpalException(
                "PartBunch::grid2par",
                "Not used: just exists because ippl::PicManager wants it that way.");
    }
 
    std::shared_ptr<VField_t<T, Dim>> getTempEField() { return this->Etmp_m; }
 
    void setTempEField(std::shared_ptr<VField_t<T, Dim>> Etmp) { this->Etmp_m = Etmp; }
 
    std::shared_ptr<VField_t<T, Dim>> getTempBField() { return this->Btmp_m; }
 
    void setTempBField(std::shared_ptr<VField_t<T, Dim>> Btmp) { this->Btmp_m = Btmp; }
 
    std::shared_ptr<AdaptBins_t> getBins() { return bins_m; }
 
    std::shared_ptr<AdaptBins_t> getBins() const { return bins_m; }
 
    void setBins(std::shared_ptr<AdaptBins_t> bins) { bins_m = bins; }
 
    void setBCHandler(std::shared_ptr<BCHandler_t> bcHandler) { bcHandler_m = bcHandler; }
 
    std::shared_ptr<BCHandler_t> getBCHandler() const { return bcHandler_m; }
 
    DataSink* getDataSink() const { return dataSink_m; }
 
    std::shared_ptr<BunchStateHandler> getBunchStateHandler() { return bunchState_m; }
    std::shared_ptr<const BunchStateHandler> getBunchStateHandler() const { return bunchState_m; }
 
    void updateMoments() { this->pcontainer_m->updateMoments(); }
 
    void calcBeamParameters();
 
    void setCharge() {
        const auto& containers = this->getParticleContainers();
        if (containers.size() != qi_m.size()) {
            throw OpalException(
                    "PartBunch::setCharge",
                    "Number of particle containers and qi values do not match.");
        }
        for (size_t i = 0; i < containers.size(); ++i) {
            containers[i]->setQ(qi_m[i]);
        }
    }
 
    void setMass() {
        const auto& containers = this->getParticleContainers();
        if (containers.size() != mi_m.size()) {
            throw OpalException(
                    "PartBunch::setMass",
                    "Number of particle containers and mi values do not match.");
        }
        for (size_t i = 0; i < containers.size(); ++i) {
            containers[i]->setM(mi_m[i]);
        }
    }
 
    double getCharge(size_t containerIndex = 0) const {
        const auto& containers = this->getParticleContainers();
        if (containers.size() != qi_m.size()) {
            throw OpalException(
                    "PartBunch::getCharge",
                    "Number of particle containers and qi values do not match.");
        }
        if (containerIndex >= containers.size()) {
            throw OpalException("PartBunch::getCharge", "Container index out of range.");
        }
        return qi_m[containerIndex] * containers[containerIndex]->getTotalNum();
    }
 
    double getChargePerParticle(size_t containerIndex = 0) const {
        if (containerIndex >= qi_m.size()) {
            throw OpalException("PartBunch::getChargePerParticle", "Container index out of range.");
        }
        return qi_m[containerIndex];
    }
 
    double getMassPerParticle(size_t containerIndex = 0) const {
        if (containerIndex >= mi_m.size()) {
            throw OpalException("PartBunch::getMassPerParticle", "Container index out of range.");
        }
        return mi_m[containerIndex];
    }
 
    double getQ(size_t containerIndex = 0) const { return this->getCharge(containerIndex); }
 
    double getM(size_t containerIndex = 0) const {
        const auto& containers = this->getParticleContainers();
        if (containers.size() != mi_m.size()) {
            throw OpalException(
                    "PartBunch::getM", "Number of particle containers and mi values do not match.");
        }
        if (containerIndex >= containers.size()) {
            throw OpalException("PartBunch::getM", "Container index out of range.");
        }
        return mi_m[containerIndex] * containers[containerIndex]->getTotalNum();
    }
 
    double getTotalCharge() const {
        const auto& containers = this->getParticleContainers();
        if (containers.size() != qi_m.size()) {
            throw OpalException(
                    "PartBunch::getTotalCharge",
                    "Number of particle containers and qi values do not match.");
        }
        double charge = 0.0;
        for (size_t i = 0; i < containers.size(); ++i) {
            charge += qi_m[i] * containers[i]->getTotalNum();
        }
        return charge;
    }
 
    double getTotalMass() const {
        const auto& containers = this->getParticleContainers();
        if (containers.size() != mi_m.size()) {
            throw OpalException(
                    "PartBunch::getTotalMass",
                    "Number of particle containers and mi values do not match.");
        }
        double mass = 0.0;
        for (size_t i = 0; i < containers.size(); ++i) {
            mass += mi_m[i] * containers[i]->getTotalNum();
        }
        return mass;
    }
 
    double getdE() {
        return this->pcontainer_m->getStdKineticEnergy();  // Unit: MeV
    }
 
    const PartData* getReference() const { return reference_m; }
 
    void setReference(const PartData* ref) {
        reference_m = ref;
        if (reference_m && this->pcontainer_m) {
            // Ensure mean/std kinetic energy in DistributionMoments are computed using reference
            // mass. PartData mass is stored in eV; DistributionMoments expects GeV for its energy
            // computation.
            this->pcontainer_m->setEnergyReferenceMass(reference_m->getM() * Units::eV2GeV, true);
        }
    }
 
    void gatherLoadBalanceStatistics();
 
    size_t getLoadBalance(int p) { return globalPartPerNode_m[p]; }
 
    size_t calcNumPartsOutside(Vector_t<double, Dim> /*x*/) {
        *gmsg << "not implemented:: file: " << __FILE__ << " line: " << __LINE__
              << " function: " << __func__ << endl;
        return 0;
    }
 
    void calcLineDensity(
            unsigned int /*nBins*/, std::vector<double>& /*lineDensity*/,
            std::pair<double, double>& /*meshInfo*/) {
        *gmsg << "not implemented:: file: " << __FILE__ << " line: " << __LINE__
              << " function: " << __func__ << endl;
    }
 
    Vector_t<double, Dim> getEExtrema() {
        *gmsg << "not implemented:: file: " << __FILE__ << " line: " << __LINE__
              << " function: " << __func__ << endl;
        return Vector_t<double, Dim>(0);
    }
 
    void computeSelfFields();
 
    void dumpBinConfig(bool preMerge);
 
    Inform& print(Inform& os);
 
    bool hasFieldSolver() const { return this->fsolver_m != nullptr; }
 
    BinnedFieldSolver_t* getFieldSolver();
 
    const BinnedFieldSolver_t* getFieldSolver() const;
 
    std::string getFieldSolverType();
 
    bool hasBinning() const { return this->bins_m != nullptr; }
 
    int getCurrentNBins() const {
        if (!hasBinning()) {
            return 1;
        }
 
        int ret_bins = static_cast<int>(bins_m->getCurrentBinCount());
        // If the number of bins is the same as the maximum number of bins, we haven't merged bins
        // yet (likely because the simulation is too empty)
        if (ret_bins == this->getBins()->getMaxBinCount()) {
            Inform m("PartBunch::getCurrentNBins");
            m << level4
              << "WARNING: Number of bins is the same as the maximum number of bins, we haven't "
                 "merged bins yet (likely because the simulation is too empty). Returning 1. If "
                 "that is not the case, check e.g. binning parameters."
              << endl;
            return 1;
        } else {
            return ret_bins;
        }
    }
 
    double calcMeanPhi() {
        *gmsg << "not implemented:: file: " << __FILE__ << " line: " << __LINE__
              << " function: " << __func__ << endl;
        return 0.0;
    }
 
    std::string getParticleName(size_t i) const {
        if (i >= particleNames_m.size()) {
            return "";
        }
        return particleNames_m[i];
    }
 
    Vector_t<double, Dim> R(size_t) {
        throw OpalException(
                "PartBunch::R",
                "Not implemented: shouldn't be called, since this is not the correct way to access "
                "particle positions.");
        return Vector_t<double, Dim>(0.0);
    }
 
    Vector_t<double, Dim> P(size_t) {
        throw OpalException(
                "PartBunch::P",
                "Not implemented: shouldn't be called, since this is not the correct way to access "
                "particle momenta.");
        return Vector_t<double, Dim>(0.0);
    }
 
    void get_bounds(Vector_t<double, Dim>& rmin, Vector_t<double, Dim>& rmax) {
        rmin = rmin_m;
        rmax = rmax_m;
    }
 
    void setdT(double dt) { dt_m = dt; }
 
    double getdT() const { return dt_m; }
 
    void setT(double t) { t_m = t; }
 
    void incrementT() { t_m += dt_m; }
 
    double getT() const { return t_m; }
 
    Vector_t<double, Dim> get_origin() const { return rmin_m; }
    Vector_t<double, Dim> get_maxExtent() const { return rmax_m; }
 
    Vector_t<double, Dim> get_halo() const {
        *gmsg << "not implemented:: file: " << __FILE__ << " line: " << __LINE__
              << " function: " << __func__ << endl;
        return Vector_t<double, Dim>(0.0);
    }
 
    Vector_t<double, Dim> get_hr() const { return hr_m; }
 
    void setGlobalTrackStep(long long n) { globalTrackStep_m = n; }
 
    long long getGlobalTrackStep() const { return globalTrackStep_m; }
 
    void incTrackSteps() { globalTrackStep_m++; }
 
    void do_binaryRepart();
 
    double get_rmsDensity() const { return rmsDensity_m; }
};
 
extern template class PartBunch<double, 3>;  
 
#endif  // PARTBUNCH_H