Distribution.cpp 168 KB
Newer Older
1
// Distribution class
gsell's avatar
gsell committed
2
//
3 4 5
// Copyright (c) 2008-2020
// Paul Scherrer Institut, Villigen PSI, Switzerland
// All rights reserved.
gsell's avatar
gsell committed
6
//
7
// OPAL is licensed under GNU GPL version 3.
kraus's avatar
kraus committed
8 9

#include "Distribution/Distribution.h"
10
#include "Distribution/SigmaGenerator.h"
11
#include "AbsBeamline/SpecificElementVisitor.h"
12

13 14 15 16 17 18
#include <cmath>
#include <cfloat>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>
19
#include <numeric>
20
#include <limits>
21

22 23 24 25
// IPPL
#include "DataSource/DataConnect.h"
#include "Utility/IpplTimings.h"

26
#include "AbstractObjects/Expressions.h"
27
#include "Utilities/Options.h"
28 29
#include "AbstractObjects/OpalData.h"
#include "Algorithms/PartBins.h"
frey_m's avatar
frey_m committed
30
#include "Algorithms/PartBunchBase.h"
31
#include "Algorithms/bet/EnvelopeBunch.h"
32
#include "Structure/Beam.h"
gsell's avatar
gsell committed
33 34 35
#include "Algorithms/PartBinsCyc.h"
#include "BasicActions/Option.h"
#include "Distribution/LaserProfile.h"
36
#include "Elements/OpalBeamline.h"
37
#include "AbstractObjects/BeamSequence.h"
38 39
#include "Structure/H5PartWrapper.h"
#include "Structure/H5PartWrapperForPC.h"
40
#include "Utilities/Util.h"
41
#include "Utilities/EarlyLeaveException.h"
42

43 44
#include <gsl/gsl_histogram.h>
#include <gsl/gsl_linalg.h>
45
#include <gsl/gsl_randist.h>
46
#include <gsl/gsl_rng.h>
47
#include <gsl/gsl_sf_erf.h>
gsell's avatar
gsell committed
48

49
#include <sys/time.h>
50

kraus's avatar
kraus committed
51
#include <boost/regex.hpp>
52
#include <boost/numeric/odeint/stepper/runge_kutta4.hpp>
kraus's avatar
kraus committed
53

gsell's avatar
gsell committed
54 55
extern Inform *gmsg;

56
constexpr double SMALLESTCUTOFF = 1e-12;
kraus's avatar
kraus committed
57

58 59 60 61 62 63
/* Increase tEmission_m by twice this percentage
 * to ensure that no particles fall on the leading edge of
 * the first emission time step or the trailing edge of the last
 * emission time step. */
const double Distribution::percentTEmission_m = 0.0005;

snuverink_j's avatar
snuverink_j committed
64 65 66 67 68 69 70
namespace {
    SymTenzor<double, 6> getUnit6x6() {
        SymTenzor<double, 6> unit6x6;
        for (unsigned int i = 0; i < 6u; ++ i) {
            unit6x6(i,i) = 1.0;
        }
        return unit6x6;
71 72 73
    }
}

gsell's avatar
gsell committed
74 75 76 77 78
//
// Class Distribution
// ------------------------------------------------------------------------

Distribution::Distribution():
79
    Definition( Attrib::Legacy::Distribution::SIZE, "DISTRIBUTION",
80
                "The DISTRIBUTION statement defines data for the 6D particle distribution."),
81
    distrTypeT_m(DistrTypeT::NODIST),
82
    numberOfDistributions_m(1),
83 84 85 86 87
    emitting_m(false),
    emissionModel_m(EmissionModelT::NONE),
    tEmission_m(0.0),
    tBin_m(0.0),
    currentEmissionTime_m(0.0),
kraus's avatar
kraus committed
88
    currentEnergyBin_m(1),
89
    currentSampleBin_m(0),
90 91 92 93
    numberOfEnergyBins_m(0),
    numberOfSampleBins_m(0),
    energyBins_m(NULL),
    energyBinHist_m(NULL),
94
    randGen_m(NULL),
95
    pTotThermal_m(0.0),
kraus's avatar
kraus committed
96
    pmean_m(0.0),
97 98 99 100 101
    cathodeWorkFunc_m(0.0),
    laserEnergy_m(0.0),
    cathodeFermiEnergy_m(0.0),
    cathodeTemp_m(0.0),
    emitEnergyUpperLimit_m(0.0),
102 103
    totalNumberParticles_m(0),
    totalNumberEmittedParticles_m(0),
104
    avrgpz_m(0.0),
105 106 107 108 109
    inputMoUnits_m(InputMomentumUnitsT::NONE),
    sigmaTRise_m(0.0),
    sigmaTFall_m(0.0),
    tPulseLengthFWHM_m(0.0),
    correlationMatrix_m(getUnit6x6()),
110 111
    sepPeaks_m(0.0),
    nPeaks_m(1),
112 113 114 115 116
    laserProfileFileName_m(""),
    laserImageName_m(""),
    laserIntensityCut_m(0.0),
    laserProfile_m(NULL),
    I_m(0.0),
117
    E_m(0.0)
118
{
119
    setAttributes();
120 121 122

    Distribution *defaultDistribution = clone("UNNAMED_Distribution");
    defaultDistribution->builtin = true;
gsell's avatar
gsell committed
123 124

    try {
125
        OpalData::getInstance()->define(defaultDistribution);
gsell's avatar
gsell committed
126
    } catch(...) {
127 128 129
        delete defaultDistribution;
    }

130 131
    gsl_rng_env_setup();
    randGen_m = gsl_rng_alloc(gsl_rng_default);
gsell's avatar
gsell committed
132 133 134 135 136 137 138
}
/**
 *
 *
 * @param name
 * @param parent
 */
139
Distribution::Distribution(const std::string &name, Distribution *parent):
gsell's avatar
gsell committed
140
    Definition(name, parent),
141 142
    distT_m(parent->distT_m),
    distrTypeT_m(DistrTypeT::NODIST),
143
    numberOfDistributions_m(parent->numberOfDistributions_m),
144 145 146 147 148
    emitting_m(parent->emitting_m),
    particleRefData_m(parent->particleRefData_m),
    addedDistributions_m(parent->addedDistributions_m),
    particlesPerDist_m(parent->particlesPerDist_m),
    emissionModel_m(parent->emissionModel_m),
gsell's avatar
gsell committed
149
    tEmission_m(parent->tEmission_m),
150 151
    tBin_m(parent->tBin_m),
    currentEmissionTime_m(parent->currentEmissionTime_m),
152
    currentEnergyBin_m(parent->currentEnergyBin_m),
153 154 155 156 157
    currentSampleBin_m(parent->currentSampleBin_m),
    numberOfEnergyBins_m(parent->numberOfEnergyBins_m),
    numberOfSampleBins_m(parent->numberOfSampleBins_m),
    energyBins_m(NULL),
    energyBinHist_m(NULL),
158
    randGen_m(NULL),
159
    pTotThermal_m(parent->pTotThermal_m),
kraus's avatar
kraus committed
160
    pmean_m(parent->pmean_m),
161 162 163 164 165
    cathodeWorkFunc_m(parent->cathodeWorkFunc_m),
    laserEnergy_m(parent->laserEnergy_m),
    cathodeFermiEnergy_m(parent->cathodeFermiEnergy_m),
    cathodeTemp_m(parent->cathodeTemp_m),
    emitEnergyUpperLimit_m(parent->emitEnergyUpperLimit_m),
166 167
    totalNumberParticles_m(parent->totalNumberParticles_m),
    totalNumberEmittedParticles_m(parent->totalNumberEmittedParticles_m),
168 169 170 171 172 173 174 175 176 177 178 179
    xDist_m(parent->xDist_m),
    pxDist_m(parent->pxDist_m),
    yDist_m(parent->yDist_m),
    pyDist_m(parent->pyDist_m),
    tOrZDist_m(parent->tOrZDist_m),
    pzDist_m(parent->pzDist_m),
    xWrite_m(parent->xWrite_m),
    pxWrite_m(parent->pxWrite_m),
    yWrite_m(parent->yWrite_m),
    pyWrite_m(parent->pyWrite_m),
    tOrZWrite_m(parent->tOrZWrite_m),
    pzWrite_m(parent->pzWrite_m),
180
    avrgpz_m(parent->avrgpz_m),
181 182 183 184 185 186 187 188
    inputMoUnits_m(parent->inputMoUnits_m),
    sigmaTRise_m(parent->sigmaTRise_m),
    sigmaTFall_m(parent->sigmaTFall_m),
    tPulseLengthFWHM_m(parent->tPulseLengthFWHM_m),
    sigmaR_m(parent->sigmaR_m),
    sigmaP_m(parent->sigmaP_m),
    cutoffR_m(parent->cutoffR_m),
    cutoffP_m(parent->cutoffP_m),
189
    correlationMatrix_m(parent->correlationMatrix_m),
190 191
    sepPeaks_m(parent->sepPeaks_m),
    nPeaks_m(parent->nPeaks_m),
192 193 194 195
    laserProfileFileName_m(parent->laserProfileFileName_m),
    laserImageName_m(parent->laserImageName_m),
    laserIntensityCut_m(parent->laserIntensityCut_m),
    laserProfile_m(NULL),
196 197
    I_m(parent->I_m),
    E_m(parent->E_m),
198 199 200 201
    tRise_m(parent->tRise_m),
    tFall_m(parent->tFall_m),
    sigmaRise_m(parent->sigmaRise_m),
    sigmaFall_m(parent->sigmaFall_m),
202
    cutoff_m(parent->cutoff_m)
203
{
204 205
    gsl_rng_env_setup();
    randGen_m = gsl_rng_alloc(gsl_rng_default);
gsell's avatar
gsell committed
206 207 208 209
}

Distribution::~Distribution() {

210 211 212 213
    delete energyBins_m;
    gsl_histogram_free(energyBinHist_m);
    gsl_rng_free(randGen_m);
    delete laserProfile_m;
gsell's avatar
gsell committed
214
}
215

216
/*
217
  void Distribution::printSigma(SigmaGenerator<double,unsigned int>::matrix_type& M, Inform& out) {
218 219
  for (int i=0; i<M.size1(); ++i) {
  for (int j=0; j<M.size2(); ++j) {
220 221 222 223
  *gmsg  << M(i,j) << " ";
  }
  *gmsg << endl;
  }
224 225
  }
*/
gsell's avatar
gsell committed
226

227 228
/**
 * Calculate the local number of particles evenly and adjust node 0
229 230 231
 * such that n is matched exactly.
 * @param n total number of particles
 * @return n / #cores
232 233 234 235 236
 * @param
 */

size_t Distribution::getNumOfLocalParticlesToCreate(size_t n) {

237
    size_t locNumber = n / Ippl::getNodes();
238

239 240 241 242 243
    // make sure the total number is exact
    size_t remainder  = n % Ippl::getNodes();
    size_t myNode = Ippl::myNode();
    if (myNode < remainder)
        ++ locNumber;
244 245 246 247

    return locNumber;
}

248 249 250
/// Distribution can only be replaced by another distribution.
bool Distribution::canReplaceBy(Object *object) {
    return dynamic_cast<Distribution *>(object) != 0;
251 252
}

253 254 255
Distribution *Distribution::clone(const std::string &name) {
    return new Distribution(name, this);
}
256

257 258
void Distribution::execute() {
}
adelmann's avatar
adelmann committed
259

260 261
void Distribution::update() {
}
adelmann's avatar
adelmann committed
262

kraus's avatar
kraus committed
263 264 265 266 267 268 269 270 271 272 273 274
void Distribution::create(size_t &numberOfParticles, double massIneV) {

    /*
     * If Options::cZero is true than we reflect generated distribution
     * to ensure that the transverse averages are 0.0.
     *
     * For now we just cut the number of generated particles in half.
     */
    size_t numberOfLocalParticles = numberOfParticles;
    if (Options::cZero && distrTypeT_m != DistrTypeT::FROMFILE) {
        numberOfLocalParticles = (numberOfParticles + 1) / 2;
    }
adelmann's avatar
adelmann committed
275

276
    size_t mySeed = Options::seed;
adelmann's avatar
adelmann committed
277

278
    if (Options::seed == -1) {
Andreas Adelmann's avatar
Andreas Adelmann committed
279 280
        struct timeval tv;
        gettimeofday(&tv,0);
281
        mySeed = tv.tv_sec + tv.tv_usec;
Andreas Adelmann's avatar
Andreas Adelmann committed
282
    }
283

284
    if (Attributes::getBool(itsAttr[Attrib::Distribution::SCALABLE])) {
kraus's avatar
kraus committed
285
        numberOfLocalParticles = getNumOfLocalParticlesToCreate(numberOfLocalParticles);
286 287 288 289 290 291
        *gmsg << level2 << "* Generation of distribution with seed = " << mySeed << " + core_id\n"
              << "* is scalable with number of particles and cores." << endl;
        mySeed += Ippl::myNode();
    } else {
        *gmsg << level2 << "* Generation of distribution with seed = " << mySeed << "\n"
              << "* isn't scalable with number of particles and cores." << endl;
Andreas Adelmann's avatar
Andreas Adelmann committed
292 293
    }

294 295
    gsl_rng_set(randGen_m, mySeed);

296
    switch (distrTypeT_m) {
adelmann's avatar
adelmann committed
297

298
    case DistrTypeT::MATCHEDGAUSS:
kraus's avatar
kraus committed
299
        createMatchedGaussDistribution(numberOfLocalParticles, massIneV);
300
        break;
301
    case DistrTypeT::FROMFILE:
kraus's avatar
kraus committed
302
        createDistributionFromFile(numberOfParticles, massIneV);
gsell's avatar
gsell committed
303
        break;
304
    case DistrTypeT::GAUSS:
kraus's avatar
kraus committed
305
        createDistributionGauss(numberOfLocalParticles, massIneV);
gsell's avatar
gsell committed
306
        break;
307
    case DistrTypeT::BINOMIAL:
kraus's avatar
kraus committed
308
        createDistributionBinomial(numberOfLocalParticles, massIneV);
gsell's avatar
gsell committed
309
        break;
310 311 312
    case DistrTypeT::FLATTOP:
    case DistrTypeT::GUNGAUSSFLATTOPTH:
    case DistrTypeT::ASTRAFLATTOPTH:
kraus's avatar
kraus committed
313
        createDistributionFlattop(numberOfLocalParticles, massIneV);
gsell's avatar
gsell committed
314
        break;
315 316 317
    case DistrTypeT::MULTIGAUSS:
        createDistributionMultiGauss(numberOfLocalParticles, massIneV);
        break;
318 319 320 321
    default:
        INFOMSG("Distribution unknown." << endl;);
        break;
    }
kraus's avatar
kraus committed
322 323

    if (emitting_m) {
324

kraus's avatar
kraus committed
325
        unsigned int numAdditionalRNsPerParticle = 0;
326
        if (emissionModel_m == EmissionModelT::ASTRA ||
327 328
            distrTypeT_m == DistrTypeT::ASTRAFLATTOPTH ||
            distrTypeT_m == DistrTypeT::GUNGAUSSFLATTOPTH) {
329

kraus's avatar
kraus committed
330
            numAdditionalRNsPerParticle = 2;
331
        } else if (emissionModel_m == EmissionModelT::NONEQUIL) {
332 333 334 335 336
            if (Options::cZero) {
                numAdditionalRNsPerParticle = 40;
            } else {
                numAdditionalRNsPerParticle = 20;
            }
kraus's avatar
kraus committed
337 338
        }

339 340 341
        int saveProcessor = -1;
        const int myNode = Ippl::myNode();
        const int numNodes = Ippl::getNodes();
342
        const bool scalable = Attributes::getBool(itsAttr[Attrib::Distribution::SCALABLE]);
343

kraus's avatar
kraus committed
344
        for (size_t partIndex = 0; partIndex < numberOfLocalParticles; ++ partIndex) {
345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361

            // Save to each processor in turn.
            ++ saveProcessor;
            if (saveProcessor >= numNodes)
                saveProcessor = 0;

            if (scalable || myNode == saveProcessor) {
                std::vector<double> rns;
                for (unsigned int i = 0; i < numAdditionalRNsPerParticle; ++ i) {
                    double x = gsl_rng_uniform(randGen_m);
                    rns.push_back(x);
                }
                additionalRNs_m.push_back(rns);
            } else {
                for (unsigned int i = 0; i < numAdditionalRNsPerParticle; ++ i) {
                    gsl_rng_uniform(randGen_m);
                }
kraus's avatar
kraus committed
362 363 364 365
            }
        }
    }

366 367
    // Scale coordinates according to distribution input.
    scaleDistCoordinates();
368

kraus's avatar
kraus committed
369 370 371
    // Now reflect particles if Options::cZero is true
    reflectDistribution(numberOfLocalParticles);

372
    adjustPhaseSpace(massIneV);
373

374
    if (Options::seed != -1)
Andreas Adelmann's avatar
Andreas Adelmann committed
375
        Options::seed = gsl_rng_uniform_int(randGen_m, gsl_rng_max(randGen_m));
kraus's avatar
kraus committed
376

377 378 379 380 381
    if (particlesPerDist_m.size() == 0) {
        particlesPerDist_m.push_back(tOrZDist_m.size());
    } else {
        particlesPerDist_m[0] = tOrZDist_m.size();
    }
382
}
gsell's avatar
gsell committed
383

384
void Distribution::doRestartOpalT(PartBunchBase<double, 3> *beam, size_t /*Np*/, int restartStep, H5PartWrapper *dataSource) {
385

frey_m's avatar
frey_m committed
386
    IpplTimings::startTimer(beam->distrReload_m);
387

388 389
    long numParticles = dataSource->getNumParticles();
    size_t numParticlesPerNode = numParticles / Ippl::getNodes();
gsell's avatar
gsell committed
390

391 392 393 394
    size_t firstParticle = numParticlesPerNode * Ippl::myNode();
    size_t lastParticle = firstParticle + numParticlesPerNode - 1;
    if (Ippl::myNode() == Ippl::getNodes() - 1)
        lastParticle = numParticles - 1;
adelmann's avatar
adelmann committed
395
    OpalData::getInstance()->addProblemCharacteristicValue("NP", numParticles);
396

397
    numParticles = lastParticle - firstParticle + 1;
398
    PAssert_GE(numParticles, 0);
399

frey_m's avatar
frey_m committed
400
    beam->create(numParticles);
401

402 403
    dataSource->readHeader();
    dataSource->readStep(beam, firstParticle, lastParticle);
404

frey_m's avatar
frey_m committed
405
    beam->boundp();
406

frey_m's avatar
frey_m committed
407 408 409
    double actualT = beam->getT();
    long long ltstep = beam->getLocalTrackStep();
    long long gtstep = beam->getGlobalTrackStep();
410

frey_m's avatar
frey_m committed
411
    IpplTimings::stopTimer(beam->distrReload_m);
gsell's avatar
gsell committed
412

frey_m's avatar
frey_m committed
413
    *gmsg << "Total number of particles in the h5 file= " << beam->getTotalNum() << "\n"
414 415 416
          << "Global step= " << gtstep << "; Local step= " << ltstep << "; "
          << "restart step= " << restartStep << "\n"
          << "time of restart= " << actualT << "; phishift= " << OpalData::getInstance()->getGlobalPhaseShift() << endl;
gsell's avatar
gsell committed
417 418
}

frey_m's avatar
frey_m committed
419
void Distribution::doRestartOpalCycl(PartBunchBase<double, 3> *beam,
420 421
                                     size_t /*Np*/,
                                     int /*restartStep*/,
422 423
                                     const int specifiedNumBunch,
                                     H5PartWrapper *dataSource) {
adelmann's avatar
adelmann committed
424

425
    // h5_int64_t rc;
frey_m's avatar
frey_m committed
426
    IpplTimings::startTimer(beam->distrReload_m);
427
    INFOMSG("---------------- Start reading hdf5 file----------------" << endl);
428

429 430
    long numParticles = dataSource->getNumParticles();
    size_t numParticlesPerNode = numParticles / Ippl::getNodes();
431

432 433 434 435
    size_t firstParticle = numParticlesPerNode * Ippl::myNode();
    size_t lastParticle = firstParticle + numParticlesPerNode - 1;
    if (Ippl::myNode() == Ippl::getNodes() - 1)
        lastParticle = numParticles - 1;
adelmann's avatar
adelmann committed
436
    OpalData::getInstance()->addProblemCharacteristicValue("NP", numParticles);
437

438
    numParticles = lastParticle - firstParticle + 1;
439
    PAssert_GE(numParticles, 0);
440

frey_m's avatar
frey_m committed
441
    beam->create(numParticles);
442

443 444
    dataSource->readHeader();
    dataSource->readStep(beam, firstParticle, lastParticle);
445

frey_m's avatar
frey_m committed
446
    beam->Q = beam->getChargePerParticle();
447

frey_m's avatar
frey_m committed
448
    beam->boundp();
449

450
    double meanE = static_cast<H5PartWrapperForPC*>(dataSource)->getMeanKineticEnergy();
451

frey_m's avatar
frey_m committed
452
    const int globalN = beam->getTotalNum();
frey_m's avatar
frey_m committed
453
    INFOMSG("Restart from hdf5 format file " << OpalData::getInstance()->getRestartFileName() << endl);
454
    INFOMSG("total number of particles = " << globalN << endl);
455
    INFOMSG("* Restart Energy = " << meanE << " (MeV), Path lenght = " << beam->get_sPos() << " (m)" <<  endl);
frey_m's avatar
frey_m committed
456 457
    INFOMSG("Tracking Step since last bunch injection is " << beam->getSteptoLastInj() << endl);
    INFOMSG(beam->getNumBunch() << " Bunches(bins) exist in this file" << endl);
458

frey_m's avatar
frey_m committed
459
    double gamma = 1 + meanE / beam->getM() * 1.0e6;
460
    double beta = std::sqrt(1.0 - (1.0 / std::pow(gamma, 2)));
461

462
    INFOMSG("* Gamma = " << gamma << ", Beta = " << beta << endl);
463

464
    if (dataSource->predecessorIsSameFlavour()) {
465
        INFOMSG("Restart from hdf5 file generated by OPAL-cycl" << endl);
466
        if (specifiedNumBunch > 1) {
467
            // the allowed maximal bin number is set to 1000
frey_m's avatar
frey_m committed
468 469
            energyBins_m = new PartBinsCyc(1000, beam->getNumBunch());
            beam->setPBins(energyBins_m);
470 471
        }

472 473
    } else {
        INFOMSG("Restart from hdf5 file generated by OPAL-t" << endl);
474 475 476

        Vector_t meanR(0.0, 0.0, 0.0);
        Vector_t meanP(0.0, 0.0, 0.0);
frey_m's avatar
frey_m committed
477
        unsigned long int newLocalN = beam->getLocalNum();
478 479
        for (unsigned int i = 0; i < newLocalN; ++i) {
            for (int d = 0; d < 3; ++d) {
frey_m's avatar
frey_m committed
480 481
                meanR(d) += beam->R[i](d);
                meanP(d) += beam->P[i](d);
482 483 484 485 486 487
            }
        }
        reduce(meanR, meanR, OpAddAssign());
        meanR /= Vector_t(globalN);
        reduce(meanP, meanP, OpAddAssign());
        meanP /= Vector_t(globalN);
488
        INFOMSG("Rmean = " << meanR << "[m], Pmean=" << meanP << endl);
489

490
        for (unsigned int i = 0; i < newLocalN; ++i) {
frey_m's avatar
frey_m committed
491 492
            beam->R[i] -= meanR;
            beam->P[i] -= meanP;
493 494 495
        }
    }

496
    INFOMSG("---------------Finished reading hdf5 file---------------" << endl);
frey_m's avatar
frey_m committed
497
    IpplTimings::stopTimer(beam->distrReload_m);
498 499
}

500
void Distribution::doRestartOpalE(EnvelopeBunch *beam, size_t /*Np*/, int /*restartStep*/,
501
                                  H5PartWrapper *dataSource) {
frey_m's avatar
frey_m committed
502
    IpplTimings::startTimer(beam->distrReload_m);
503 504
    int N = dataSource->getNumParticles();
    *gmsg << "total number of slices = " << N << endl;
505

frey_m's avatar
frey_m committed
506 507 508 509
    beam->distributeSlices(N);
    beam->createBunch();
    long long starti = beam->mySliceStartOffset();
    long long endi = beam->mySliceEndOffset();
510

511 512
    dataSource->readHeader();
    dataSource->readStep(beam, starti, endi);
513

frey_m's avatar
frey_m committed
514 515
    beam->setCharge(beam->getChargePerParticle());
    IpplTimings::stopTimer(beam->distrReload_m);
516 517 518 519 520
}

Distribution *Distribution::find(const std::string &name) {
    Distribution *dist = dynamic_cast<Distribution *>(OpalData::getInstance()->find(name));

521
    if (dist == 0) {
522 523 524 525 526 527
        throw OpalException("Distribution::find()", "Distribution \"" + name + "\" not found.");
    }

    return dist;
}

528
double Distribution::getTEmission() {
529
    if (tEmission_m > 0.0) {
530 531 532
        return tEmission_m;
    }

533
    setDistType();
534

535 536 537 538
    tPulseLengthFWHM_m = Attributes::getReal(itsAttr[Attrib::Distribution::TPULSEFWHM]);
    cutoff_m = Attributes::getReal(itsAttr[Attrib::Legacy::Distribution::CUTOFF]);
    tRise_m = Attributes::getReal(itsAttr[Attrib::Distribution::TRISE]);
    tFall_m = Attributes::getReal(itsAttr[Attrib::Distribution::TFALL]);
539
    double tratio = std::sqrt(2.0 * std::log(10.0)) - std::sqrt(2.0 * std::log(10.0 / 9.0));
540 541 542 543
    sigmaRise_m = tRise_m / tratio;
    sigmaFall_m = tFall_m / tratio;

    switch(distrTypeT_m) {
544 545 546 547
    case DistrTypeT::ASTRAFLATTOPTH: {
        double a = tPulseLengthFWHM_m / 2;
        double sig = tRise_m / 2;
        double inv_erf08 = 0.906193802436823; // erfinv(0.8)
548
        double sqr2 = std::sqrt(2.);
549 550 551
        double t = a - sqr2 * sig * inv_erf08;
        double tmps = sig;
        double tmpt = t;
552
        for (int i = 0; i < 10; ++ i) {
553 554 555 556 557
            sig = (t + tRise_m - a) / (sqr2 * inv_erf08);
            t = a - 0.5 * sqr2 * (sig + tmps) * inv_erf08;
            sig = (0.5 * (t + tmpt) + tRise_m - a) / (sqr2 * inv_erf08);
            tmps = sig;
            tmpt = t;
558
        }
559 560 561 562
        tEmission_m = tPulseLengthFWHM_m + 10 * sig;
        break;
    }
    case DistrTypeT::GUNGAUSSFLATTOPTH: {
563
        tEmission_m = tPulseLengthFWHM_m + (cutoff_m - std::sqrt(2.0 * std::log(2.0))) * (sigmaRise_m + sigmaFall_m);
564 565 566 567
        break;
    }
    default:
        tEmission_m = 0.0;
568 569 570 571 572 573
    }
    return tEmission_m;
}

Inform &Distribution::printInfo(Inform &os) const {

kraus's avatar
kraus committed
574 575
    os << "\n"
       << "* ************* D I S T R I B U T I O N ********************************************" << endl;
576
    os << "* " << endl;
577
    if (OpalData::getInstance()->inRestartRun()) {
adelmann's avatar
adelmann committed
578
        os << "* In restart. Distribution read in from .h5 file." << endl;
579 580
    } else {
        if (addedDistributions_m.size() > 0)
adelmann's avatar
adelmann committed
581
            os << "* Main Distribution" << endl
582
               << "-----------------" << endl;
583 584

        if (particlesPerDist_m.empty())
585
            printDist(os, 0);
586
        else
587
            printDist(os, particlesPerDist_m.at(0));
588 589 590

        size_t distCount = 1;
        for (unsigned distIndex = 0; distIndex < addedDistributions_m.size(); distIndex++) {
591
            os << "* " << endl;
adelmann's avatar
adelmann committed
592 593
            os << "* Added Distribution #" << distCount << endl;
            os << "* ----------------------" << endl;
594
            addedDistributions_m.at(distIndex)->printDist(os, particlesPerDist_m.at(distCount));
595 596 597
            distCount++;
        }

598
        os << "* " << endl;
599
        if (numberOfEnergyBins_m > 0) {
adelmann's avatar
adelmann committed
600
            os << "* Number of energy bins    = " << numberOfEnergyBins_m << endl;
601

602
            //            if (numberOfEnergyBins_m > 1)
603
            //    printEnergyBins(os);
604 605 606
        }

        if (emitting_m) {
adelmann's avatar
adelmann committed
607 608
            os << "* Distribution is emitted. " << endl;
            os << "* Emission time            = " << tEmission_m << " [sec]" << endl;
609 610
            os << "* Time per bin             = " << tEmission_m / numberOfEnergyBins_m << " [sec]" << endl;
            os << "* Delta t during emission  = " << tBin_m / numberOfSampleBins_m << " [sec]" << endl;
611
            os << "* " << endl;
612
            printEmissionModel(os);
kraus's avatar
kraus committed
613
        } else {
adelmann's avatar
adelmann committed
614
            os << "* Distribution is injected." << endl;
kraus's avatar
kraus committed
615
        }
616
    }
617 618
    os << "* " << endl;
    os << "* **********************************************************************************" << endl;
619 620 621 622

    return os;
}

623
void Distribution::addDistributions() {
624 625 626
    /*
     * Move particle coordinates from added distributions to main distribution.
     */
adelmann's avatar
doc  
adelmann committed
627

kraus's avatar
kraus committed
628
    size_t idx = 1;
629 630
    std::vector<Distribution *>::iterator addedDistIt;
    for (addedDistIt = addedDistributions_m.begin();
kraus's avatar
kraus committed
631
         addedDistIt != addedDistributions_m.end(); ++ addedDistIt, ++ idx) {
632

kraus's avatar
kraus committed
633
        particlesPerDist_m[idx] = (*addedDistIt)->tOrZDist_m.size();
634

635 636
        for (double dist : (*addedDistIt)->getXDist()) {
            xDist_m.push_back(dist);
637
        }
638
        (*addedDistIt)->eraseXDist();
639

640 641
        for (double dist : (*addedDistIt)->getBGxDist()) {
            pxDist_m.push_back(dist);
642
        }
643
        (*addedDistIt)->eraseBGxDist();
644

645 646
        for (double dist : (*addedDistIt)->getYDist()) {
            yDist_m.push_back(dist);
647
        }
648
        (*addedDistIt)->eraseYDist();
649

650 651
        for (double dist : (*addedDistIt)->getBGyDist()) {
            pyDist_m.push_back(dist);
652
        }
653
        (*addedDistIt)->eraseBGyDist();
654

655 656
        for (double dist : (*addedDistIt)->getTOrZDist()) {
            tOrZDist_m.push_back(dist);
657
        }
658
        (*addedDistIt)->eraseTOrZDist();
659

660 661
        for (double dist : (*addedDistIt)->getBGzDist()) {
            pzDist_m.push_back(dist);
662
        }
663
        (*addedDistIt)->eraseBGzDist();
664 665 666
    }
}

667
void Distribution::applyEmissionModel(double lowEnergyLimit, double &px, double &py, double &pz, std::vector<double> &additionalRNs) {
668 669 670 671

    switch (emissionModel_m) {

    case EmissionModelT::NONE:
672
        applyEmissModelNone(pz);
673 674
        break;
    case EmissionModelT::ASTRA:
675
        applyEmissModelAstra(px, py, pz, additionalRNs);
676 677
        break;
    case EmissionModelT::NONEQUIL:
678
        applyEmissModelNonEquil(lowEnergyLimit, px, py, pz, additionalRNs);
679 680 681 682 683 684
        break;
    default:
        break;
    }
}

685
void Distribution::applyEmissModelAstra(double &px, double &py, double &pz, std::vector<double> &additionalRNs) {
686

687
    double phi = 2.0 * std::acos(std::sqrt(additionalRNs[0]));
kraus's avatar
kraus committed
688
    double theta = 2.0 * Physics::pi * additionalRNs[1];
689

690 691 692
    px = pTotThermal_m * std::sin(phi) * std::cos(theta);
    py = pTotThermal_m * std::sin(phi) * std::sin(theta);
    pz = pTotThermal_m * std::abs(std::cos(phi));
693 694 695

}

696
void Distribution::applyEmissModelNone(double &pz) {
697 698 699
    pz += pTotThermal_m;
}

700
void Distribution::applyEmissModelNonEquil(double lowEnergyLimit,
701 702
                                           double &bgx,
                                           double &bgy,
703 704
                                           double &bgz,
                                           std::vector<double> &additionalRNs) {
705 706 707 708

    // Generate emission energy.
    double energy = 0.0;
    bool allow = false;
709 710 711 712

    const double expRelativeLaserEnergy = exp(laserEnergy_m / cathodeTemp_m);
    // double energyRange = emitEnergyUpperLimit_m - lowEnergyLimit;
    unsigned int counter = 0;
713
    while (!allow) {
714 715 716 717 718 719 720
        energy = lowEnergyLimit + additionalRNs[counter++] * emitEnergyUpperLimit_m;
        double randFuncValue = additionalRNs[counter++];
        double expRelativeEnergy = exp((energy - cathodeFermiEnergy_m) / cathodeTemp_m);
        double funcValue = ((1.0
                            - 1.0 / (1.0 + expRelativeEnergy * expRelativeLaserEnergy)) /
                            (1.0 + expRelativeEnergy));

721 722
        if (randFuncValue <= funcValue)
            allow = true;
723 724 725 726 727 728 729 730 731

        if (counter == additionalRNs.size()) {
            for (unsigned int i = 0; i < counter; ++ i) {
                additionalRNs[i] = gsl_rng_uniform(randGen_m);
            }

            counter = 0;
        }
    }
732

733 734 735
    while (additionalRNs.size() - counter < 2) {
        -- counter;
        additionalRNs[counter] = gsl_rng_uniform(randGen_m);
736 737 738 739
    }

    // Compute emission angles.
    double energyInternal = energy + laserEnergy_m;
740
    double energyExternal = energy - lowEnergyLimit; // uniformly distributed (?) value between 0 and emitEnergyUpperLimit_m
741

742
    double thetaInMax = std::acos(std::sqrt((lowEnergyLimit + laserEnergy_m) / (energyInternal)));
743
    double thetaIn = additionalRNs[counter++] * thetaInMax;
744
    double sinThetaOut = std::sin(thetaIn) * std::sqrt(energyInternal / energyExternal);
745
    double phi = Physics::two_pi * additionalRNs[counter];
746 747 748

    // Compute emission momenta.
    double betaGammaExternal
749
        = std::sqrt(std::pow(energyExternal / (Physics::m_e * 1.0e9) + 1.0, 2) - 1.0);
750

751 752 753
    bgx = betaGammaExternal * sinThetaOut * std::cos(phi);
    bgy = betaGammaExternal * sinThetaOut * std::sin(phi);
    bgz = betaGammaExternal * std::sqrt(1.0 - std::pow(sinThetaOut, 2));
754 755 756

}

757
void Distribution::calcPartPerDist(size_t numberOfParticles) {
758

kraus's avatar
kraus committed
759
    if (numberOfDistributions_m == 1) {
760
        particlesPerDist_m.push_back(numberOfParticles);
kraus's avatar
kraus committed
761 762 763 764 765 766 767 768 769 770 771 772 773
        return;
    }

    std::map<unsigned int, size_t> nPartFromFiles;
    double totalWeight = 0.0;
    for (unsigned int i = 0; i <= addedDistributions_m.size(); ++ i) {
        Distribution *currDist = this;
        if (i > 0)
            currDist = addedDistributions_m[i - 1];

        if (currDist->distrTypeT_m == DistrTypeT::FROMFILE) {
            std::ifstream inputFile;
            if (Ippl::myNode() == 0) {
774
                std::string fileName = Attributes::getString(currDist->itsAttr[Attrib::Distribution::FNAME]);
kraus's avatar
kraus committed
775 776 777 778 779 780
                inputFile.open(fileName.c_str());
            }
            size_t nPart = getNumberOfParticlesInFile(inputFile);
            nPartFromFiles.insert(std::make_pair(i, nPart));
            if (nPart > numberOfParticles) {
                throw OpalException("Distribution::calcPartPerDist",
781
                                    "Number of particles is too small");
kraus's avatar
kraus committed
782 783 784 785 786
            }

            numberOfParticles -= nPart;
        } else {
            totalWeight += currDist->getWeight();
787
        }
kraus's avatar
kraus committed
788
    }
789

kraus's avatar
kraus committed
790 791 792 793 794
    size_t numberOfCommittedPart = 0;
    for (unsigned int i = 0; i <= addedDist