Distribution.cpp

Go to the documentation of this file.

//
// Class Distribution
//   This class defines the initial beam that is injected or emitted into the simulation.
//
// Copyright (c) 2008 - 2022, Paul Scherrer Institut, Villigen PSI, Switzerland
// All rights reserved
//
// This file is part of OPAL.
//
// OPAL is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// You should have received a copy of the GNU General Public License
// along with OPAL. If not, see <https://www.gnu.org/licenses/>.
//
#include "Distribution/Distribution.h"
#include "AbstractObjects/BeamSequence.h"
#include "AbstractObjects/Expressions.h"
#include "AbstractObjects/OpalData.h"
#include "BasicActions/Option.h"
#include "Distribution/LaserProfile.h"
#include "Elements/OpalBeamline.h"
#include "OPALTypes.h"
#include "Physics/Physics.h"
#include "Physics/Units.h"
#include "Structure/H5PartWrapper.h"
#include "Utilities/OpalException.h"
#include "Utilities/Options.h"
#include "Utilities/Util.h"
#include "Utility/IpplTimings.h"
 
#include "Utilities/GSLCompat.h"
#include "Utilities/Random.h"
 
#include <filesystem>
#include <regex>
 
#include <sys/time.h>
 
#include <cfloat>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
 
extern Inform* gmsg;
 
using GeneratorPool = typename Kokkos::Random_XorShift64_Pool<>;
 
using view_type = typename ippl::detail::ViewType<ippl::Vector<double, Dim>, 1>::view_type;
 
namespace DISTRIBUTION {
    enum {
        TYPE,
        FNAME,
        SIGMAX,
        SIGMAY,
        SIGMAZ,
        SIGMAPX,
        SIGMAPY,
        SIGMAPZ,
        CORR,
        CUTOFFPX,
        CUTOFFPY,
        CUTOFFPZ,
        CUTOFFX,
        CUTOFFY,
        CUTOFFLONG,
        CORRX,
        CORRY,
        CORRZ,
        CORRT,
        SIGMAT,
        TPULSEFWHM,
        TRISE,
        TFALL,
        FTOSCAMPLITUDE,
        FTOSCPERIODS,
        EMITTED,
        EMISSIONSTEPS,
        NPARTDIST,
        SIZE
    };
}
 
/*
namespace {
    matrix6x6_t getUnit6x6() {
        matrix6x6_t unit6x6(0.0);  // Initialize a 6x6 matrix with all elements as 0.0
        for (unsigned int i = 0; i < 6u; ++i) {
            unit6x6(i, i) = 1.0;  // Set diagonal elements to 1.0
        }
        return unit6x6;
    }
}
*/
 
Distribution::Distribution()
    : Definition(
              DISTRIBUTION::SIZE, "DISTRIBUTION",
              "The DISTRIBUTION statement defines data for the 6D particle distribution."),
      totalNumberParticles_m(0),
      distrTypeT_m(DistributionType::NODIST),
      avrgpz_m(0.0) {
    itsAttr[DISTRIBUTION::TYPE] = Attributes::makePredefinedString(
            "TYPE", "Distribution type.",
            {"GAUSS", "MULTIVARIATEGAUSS", "FLATTOP", "OPALFLATTOP", "FROMFILE",
             "EMITTEDFROMFILE"});
 
    itsAttr[DISTRIBUTION::FNAME] =
            Attributes::makeString("FNAME", "File for reading in 6D particle coordinates.", "");
 
    // Parameters for defining an initial distribution.
    itsAttr[DISTRIBUTION::SIGMAX]  = Attributes::makeReal("SIGMAX", "SIGMAx (m)", 0.0);
    itsAttr[DISTRIBUTION::SIGMAY]  = Attributes::makeReal("SIGMAY", "SIGMAy (m)", 0.0);
    itsAttr[DISTRIBUTION::SIGMAZ]  = Attributes::makeReal("SIGMAZ", "SIGMAz (m)", 0.0);
    itsAttr[DISTRIBUTION::SIGMAPX] = Attributes::makeReal("SIGMAPX", "SIGMApx", 0.0);
    itsAttr[DISTRIBUTION::SIGMAPY] = Attributes::makeReal("SIGMAPY", "SIGMApy", 0.0);
    itsAttr[DISTRIBUTION::SIGMAPZ] = Attributes::makeReal("SIGMAPZ", "SIGMApz", 0.0);
 
    itsAttr[DISTRIBUTION::CORR] = Attributes::makeRealArray("CORR", "r correlation");
 
    itsAttr[DISTRIBUTION::CUTOFFPX] = Attributes::makeReal(
            "CUTOFFPX", "Distribution cutoff px dimension in units of sigma.", 3.0);
    itsAttr[DISTRIBUTION::CUTOFFPY] = Attributes::makeReal(
            "CUTOFFPY", "Distribution cutoff py dimension in units of sigma.", 3.0);
    itsAttr[DISTRIBUTION::CUTOFFPZ] = Attributes::makeReal(
            "CUTOFFPZ", "Distribution cutoff pz dimension in units of sigma.", 3.0);
 
    itsAttr[DISTRIBUTION::CUTOFFX] = Attributes::makeReal(
            "CUTOFFX", "Distribution cutoff x direction in units of sigma.", 3.0);
    itsAttr[DISTRIBUTION::CUTOFFY] = Attributes::makeReal(
            "CUTOFFY", "Distribution cutoff r direction in units of sigma.", 3.0);
    itsAttr[DISTRIBUTION::CUTOFFLONG] = Attributes::makeReal(
            "CUTOFFLONG", "Distribution cutoff z or t direction in units of sigma.", 3.0);
 
    itsAttr[DISTRIBUTION::CORRX] =
            Attributes::makeReal("CORRX", "x/px correlation, (R12 in transport notation).", 0.0);
    itsAttr[DISTRIBUTION::CORRY] =
            Attributes::makeReal("CORRY", "y/py correlation, (R34 in transport notation).", 0.0);
    itsAttr[DISTRIBUTION::CORRZ] =
            Attributes::makeReal("CORRZ", "z/pz correlation, (R56 in transport notation).", 0.0);
    itsAttr[DISTRIBUTION::CORRT] =
            Attributes::makeReal("CORRT", "t/pt correlation, (R56 in transport notation).", 0.0);
 
    itsAttr[DISTRIBUTION::SIGMAT] = Attributes::makeReal("SIGMAT", "SIGMAt (m)", 0.0);
    itsAttr[DISTRIBUTION::TPULSEFWHM] =
            Attributes::makeReal("TPULSEFWHM", "Pulse FWHM for emitted distribution.", 0.0);
    itsAttr[DISTRIBUTION::TRISE] =
            Attributes::makeReal("TRISE", "Rise time for emitted distribution.", 0.0);
    itsAttr[DISTRIBUTION::TFALL] =
            Attributes::makeReal("TFALL", "Fall time for emitted distribution.", 0.0);
 
    itsAttr[DISTRIBUTION::FTOSCAMPLITUDE] = Attributes::makeReal(
            "FTOSCAMPLITUDE",
            "Amplitude of oscillations superimposed "
            "on flat top portion of emitted GAUSS "
            "distribtuion (in percent of flat top "
            "amplitude)",
            0.0);
 
    itsAttr[DISTRIBUTION::FTOSCPERIODS] = Attributes::makeReal(
            "FTOSCPERIODS",
            "Number of oscillations superimposed on "
            "flat top portion of emitted GAUSS "
            "distribution",
            0.0);
 
    itsAttr[DISTRIBUTION::EMITTED] = Attributes::makeBool(
            "EMITTED",
            "Emitted beam, from cathode, as opposed to "
            "an injected beam.",
            false);
 
    itsAttr[DISTRIBUTION::EMISSIONSTEPS] = Attributes::makeReal(
            "EMISSIONSTEPS",
            "Number of OPAL-like time steps to use during OPALFLATTOP or EMITTEDFROMFILE "
            "emission.",
            100.0);
 
    itsAttr[DISTRIBUTION::NPARTDIST] = Attributes::makeReal(
            "NPARTDIST",
            "Number of macroparticles for this DISTRIBUTION. "
            "If 0 or negative, this distribution does not specify its own count.",
            0.0);
 
    registerOwnership(AttributeHandler::STATEMENT);
}
 
Distribution::Distribution(const std::string& name, Distribution* parent)
    : Definition(name, parent) {}
 
Distribution::~Distribution() {}
 
size_t Distribution::getNumOfLocalParticlesToCreate(size_t n) {
    size_t locNumber = n / ippl::Comm->size();
 
    // make sure the total number is exact
    size_t remainder = n % ippl::Comm->size();
    size_t myNode    = ippl::Comm->rank();
    if (myNode < remainder) ++locNumber;
 
    return locNumber;
}
 
bool Distribution::canReplaceBy(Object* object) { return dynamic_cast<Distribution*>(object) != 0; }
 
Distribution* Distribution::clone(const std::string& name) { return new Distribution(name, this); }
 
void Distribution::execute() {
    setAttributes();
    update();
}
 
Distribution* Distribution::find(const std::string& name) {
    Distribution* dist = dynamic_cast<Distribution*>(OpalData::getInstance()->find(name));
 
    if (dist == 0) {
        throw OpalException("Distribution::find()", "Distribution \"" + name + "\" not found.");
    }
 
    return dist;
}
 
Inform& Distribution::printInfo(Inform& os) const {
    os << "\n"
       << "* ************* D I S T R I B U T I O N ********************************************"
       << endl;
    os << "* " << endl;
    if (OpalData::getInstance()->inRestartRun()) {
        os << "* In restart. Distribution read in from .h5 file." << endl;
    } else {
        switch (distrTypeT_m) {
            case DistributionType::GAUSS:
                printDistGauss(os);
                break;
            case DistributionType::MULTIVARIATEGAUSS:
                printDistMultiVariateGauss(os);
                break;
            case DistributionType::FLATTOP:
            case DistributionType::OPALFLATTOP:
                printDistFlatTop(os);
                break;
            case DistributionType::FROMFILE:
            case DistributionType::EMITTEDFROMFILE:
                printDistFromFile(os);
                break;
            default:
                throw OpalException("Distribution Param", "Unknown \"TYPE\" of \"DISTRIBUTION\"");
        }
        os << "* " << endl;
        os << "* Distribution is injected." << endl;
    }
    os << "* " << endl;
    os << "* **********************************************************************************"
       << endl;
 
    return os;
}
 
void Distribution::setAvrgPz(double avrgpz) { avrgpz_m = avrgpz; }
 
void Distribution::setTEmission(double tEmission) { tEmission_m = tEmission; }
 
double Distribution::getTEmission() const { return tEmission_m; }
 
size_t Distribution::getEmissionSteps() const {
    const double raw = Attributes::getReal(itsAttr[DISTRIBUTION::EMISSIONSTEPS]);
    if (raw <= 0.0) {
        return 100;
    }
    return static_cast<size_t>(std::ceil(raw));
}
 
void Distribution::setDistParametersGauss() {
    /*
     * Set distribution parameters. Do all the necessary checks depending
     * on the input attributes.
     * In case of DistributionType::MATCHEDGAUSS we only need to set the cutoff parameters
     */
 
    cutoffR_m = 3.;
    cutoffP_m = 3.;
    /*
    cutoffP_m = ippl::Vector<double, 3>(Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFPX]),
                         Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFPY]),
                         Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFPZ]));
 
 
    cutoffR_m = ippl::Vector<double, 3>(Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFX]),
                         Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFY]),
                         Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFLONG]));
    */
 
    // if (std::abs(Attributes::getReal(itsAttr[Attrib::Distribution::SIGMAR])) > 0.0) {
    //     cutoffR_m[0] = Attributes::getReal(itsAttr[Attrib::Distribution::CUTOFFR]);
    //     cutoffR_m[1] = Attributes::getReal(itsAttr[Attrib::Distribution::CUTOFFR]);
    // }
 
    setSigmaR_m();
    setSigmaP_m();
 
    avrgpz_m = 0.0;
}
 
void Distribution::setDistParametersMultiVariateGauss() {
    cutoffR_m = 3.;
    cutoffP_m = 3.;
 
    // initialize the covariance matrix to identity
    for (unsigned int i = 0; i < 6; ++i) {
        for (unsigned int j = 0; j < 6; ++j) {
            if (i == j)
                correlationMatrix_m[i][j] = 1.0;
            else
                correlationMatrix_m[i][j] = 0.0;
        }
    }
 
    // set diagonal elements first
    setSigmaR_m();
    setSigmaP_m();
 
    for (unsigned int i = 0; i < 3; ++i) {
        correlationMatrix_m[2 * i][2 * i]         = sigmaR_m[i] * sigmaR_m[i];
        correlationMatrix_m[2 * i + 1][2 * i + 1] = sigmaP_m[i] * sigmaP_m[i];
    }
 
    std::vector<double> cr = Attributes::getRealArray(itsAttr[DISTRIBUTION::CORR]);
 
    if (!cr.empty()) {
        // read off-diagonal correlation matrix from input file
        if (cr.size() == 15) {
            *gmsg << "* Use r to specify correlations" << endl;
            unsigned int k = 0;
            for (unsigned int i = 0; i < 5; ++i) {
                for (unsigned int j = i + 1; j < 6; ++j, ++k) {
                    correlationMatrix_m[j][i] = cr.at(k) * cr.at(k);
                    correlationMatrix_m[i][j] = correlationMatrix_m[j][i];  // impose symmetry
                }
            }
        } else {
            throw OpalException(
                    "Distribution::SetDistParametersGauss",
                    "Inconsistent set of correlations specified, check manual");
        }
    }
 
    avrgpz_m = 0.0;
}
 
void Distribution::setDistParametersFlatTop() {
    // avrgpz_m is not reset here; TrackRun sets it from BEAM (getMomentum()/getMass() = beta*gamma)
    // so that initial and emitted particles get the correct reference momentum.
 
    cutoffR_m = 3.;
    cutoffP_m = 3.;
 
    // set diagonal elements first
    setSigmaR_m();
    setSigmaP_m();
 
    // initialize the covariance matrix to identity
    for (unsigned int i = 0; i < 6; ++i) {
        for (unsigned int j = 0; j < 6; ++j) {
            if (i == j)
                correlationMatrix_m[i][j] = 1.0;
            else
                correlationMatrix_m[i][j] = 0.0;
        }
    }
 
    cutoffR_m = ippl::Vector<double, 3>(
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFX])),
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFY])),
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::CUTOFFLONG])));
 
    correlationMatrix_m[1][0] = Attributes::getReal(itsAttr[DISTRIBUTION::CORRX]);
    correlationMatrix_m[3][2] = Attributes::getReal(itsAttr[DISTRIBUTION::CORRY]);
    correlationMatrix_m[5][4] = Attributes::getReal(itsAttr[DISTRIBUTION::CORRT]);
 
    if (Attributes::getReal(itsAttr[DISTRIBUTION::CORRZ]) != 0.0)
        correlationMatrix_m[5][4] = Attributes::getReal(itsAttr[DISTRIBUTION::CORRZ]);
 
    emitting_m = Attributes::getBool(itsAttr[DISTRIBUTION::EMITTED]);
 
    if (emitting_m) {
        sigmaR_m[2] = 0.0;
 
        sigmaTRise_m       = std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAT]));
        sigmaTFall_m       = std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAT]));
        tPulseLengthFWHM_m = std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TPULSEFWHM]));
 
        FTOSCAmplitude_m = std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::FTOSCAMPLITUDE]));
        FTOSCPeriods_m   = std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::FTOSCPERIODS]));
 
        // If TRISE and TFALL are defined > 0.0 then these attributes
        // override SIGMAT.
        //
        if (std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TRISE])) > 0.0
            || std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TFALL])) > 0.0) {
            double timeRatio =
                    std::sqrt(2.0 * std::log(10.0)) - std::sqrt(2.0 * std::log(10.0 / 9.0));
 
            if (std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TRISE])) > 0.0)
                sigmaTRise_m =
                        std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TRISE])) / timeRatio;
 
            if (std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TFALL])) > 0.0)
                sigmaTFall_m =
                        std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::TFALL])) / timeRatio;
        }
 
        // For an emitted beam, the longitudinal cutoff >= 0.
        cutoffR_m[2] = std::abs(cutoffR_m[2]);
    }
}
 
void Distribution::printDistGauss(Inform& os) const {
    os << "* Distribution type: GAUSS" << endl;
    os << "* " << endl;
    os << "* SIGMAX     = " << sigmaR_m[0] << " [m]" << endl;
    os << "* SIGMAY     = " << sigmaR_m[1] << " [m]" << endl;
    os << "* SIGMAZ     = " << sigmaR_m[2] << " [m]" << endl;
    os << "* SIGMAPX    = " << sigmaP_m[0] << " [Beta Gamma]" << endl;
    os << "* SIGMAPY    = " << sigmaP_m[1] << " [Beta Gamma]" << endl;
    os << "* SIGMAPZ    = " << sigmaP_m[2] << " [Beta Gamma]" << endl;
}
 
void Distribution::printDistMultiVariateGauss(Inform& os) const {
    os << "* Distribution type: MULTIVARIATEGAUSS" << endl;
    os << "* " << endl;
    os << "* SIGMAX     = " << sigmaR_m[0] << " [m]" << endl;
    os << "* SIGMAY     = " << sigmaR_m[1] << " [m]" << endl;
    os << "* SIGMAZ     = " << sigmaR_m[2] << " [m]" << endl;
    os << "* SIGMAPX    = " << sigmaP_m[0] << " [Beta Gamma]" << endl;
    os << "* SIGMAPY    = " << sigmaP_m[1] << " [Beta Gamma]" << endl;
    os << "* SIGMAPZ    = " << sigmaP_m[2] << " [Beta Gamma]" << endl;
 
    os << "* input cov matrix = ";
    for (unsigned int i = 0; i < 6; ++i) {
        for (unsigned int j = 0; j < 6; ++j) {
            os << correlationMatrix_m[i][j] << " ";
        }
        os << endl << "                     ";
    }
}
 
void Distribution::printDistFlatTop(Inform& os) const {
    os << "* Distribution type: " << distT_m << endl;
    os << "* " << endl;
    os << "* SIGMAX     = " << sigmaR_m[0] << " [m]" << endl;
    os << "* SIGMAY     = " << sigmaR_m[1] << " [m]" << endl;
 
    if (emitting_m) {
        os << "* Sigma Time Rise               = " << sigmaTRise_m << " [sec]" << endl;
        os << "* TPULSEFWHM                    = " << tPulseLengthFWHM_m << " [sec]" << endl;
        os << "* Sigma Time Fall               = " << sigmaTFall_m << " [sec]" << endl;
        os << "* Longitudinal cutoff           = " << cutoffR_m[2] << " [units of Sigma Time]"
           << endl;
        // os << "* Flat top modulation amplitude = "
        //    << Attributes::getReal(itsAttr[DISTRIBUTION::FTOSCAMPLITUDE])
        //    << " [Percent of distribution amplitude]" << endl;
        // os << "* Flat top modulation periods   = "
        //    << std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::FTOSCPERIODS]))
        //    << endl;
    } else {
        os << "* SIGMAZ                        = " << sigmaR_m[2] << " [m]" << endl;
    }
}
 
void Distribution::printDistFromFile(Inform& os) const {
    os << "* Distribution type: " << distT_m << endl;
    os << "* " << endl;
    std::string fname = getFilename();
    if (!fname.empty()) {
        os << "* Filename: " << fname << endl;
    } else {
        os << "* Filename: (not set)" << endl;
    }
}
 
std::string Distribution::getFilename() const {
    return Attributes::getString(itsAttr[DISTRIBUTION::FNAME]);
}
 
void Distribution::setAttributes() {
    // Set distribution type and shape-related parameters.
    setDist();
 
    // Cache per-distribution particle count, if provided.
    // A value <= 0 means \"unspecified\" and will be handled by TrackRun/BEAM.
    const double npartDist = Attributes::getReal(itsAttr[DISTRIBUTION::NPARTDIST]);
    totalNumberParticles_m =
            npartDist > 0.0 ? static_cast<size_t>(npartDist) : static_cast<size_t>(0);
}
 
void Distribution::setDist() {
    // set distribution type
    setDistType();
    // set distribution parameters
    switch (distrTypeT_m) {
        case DistributionType::GAUSS:
            setDistParametersGauss();
            break;
        case DistributionType::MULTIVARIATEGAUSS:
            setDistParametersMultiVariateGauss();
            break;
        case DistributionType::FLATTOP:
        case DistributionType::OPALFLATTOP:
            setDistParametersFlatTop();
            break;
        case DistributionType::FROMFILE:
        case DistributionType::EMITTEDFROMFILE:
            // File-based distributions read their records in the sampler class.
            break;
        default:
            throw OpalException("Distribution Param", "Unknown \"TYPE\" of \"DISTRIBUTION\"");
    }
}
 
void Distribution::setDistType() {
    static const std::map<std::string, DistributionType> typeStringToDistType_s = {
            {"NODIST", DistributionType::NODIST},
            {"GAUSS", DistributionType::GAUSS},
            {"MULTIVARIATEGAUSS", DistributionType::MULTIVARIATEGAUSS},
            {"FLATTOP", DistributionType::FLATTOP},
            {"OPALFLATTOP", DistributionType::OPALFLATTOP},
            {"FROMFILE", DistributionType::FROMFILE},
            {"EMITTEDFROMFILE", DistributionType::EMITTEDFROMFILE}};
 
    distT_m = Attributes::getString(itsAttr[DISTRIBUTION::TYPE]);
 
    if (distT_m.empty()) {
        throw OpalException(
                "Distribution::setDistType",
                "The attribute \"TYPE\" isn't set for the \"DISTRIBUTION\"!");
    } else {
        distrTypeT_m = typeStringToDistType_s.at(distT_m);
    }
}
 
void Distribution::setSigmaR_m() {
    sigmaR_m = ippl::Vector<double, 3>(
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAX])),
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAY])),
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAZ])));
}
 
void Distribution::setSigmaP_m() {
    sigmaP_m = ippl::Vector<double, 3>(
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAPX])),
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAPY])),
            std::abs(Attributes::getReal(itsAttr[DISTRIBUTION::SIGMAPZ])));
}