Distribution.cpp 4.92 KB
Newer Older
frey_m's avatar
frey_m committed
1 2 3
#include "Distribution.h"

#include "H5Reader.h"
4
#include <fstream>
frey_m's avatar
frey_m committed
5 6 7 8 9 10


// ----------------------------------------------------------------------------
// PUBLIC MEMBER FUNCTIONS
// ----------------------------------------------------------------------------

11 12 13 14 15 16 17 18 19 20 21 22
Distribution::Distribution():
    x_m(0),
    y_m(0),
    z_m(0),
    px_m(0),
    py_m(0),
    pz_m(0),
    nloc_m(0),
    ntot_m(0)
{ }


frey_m's avatar
frey_m committed
23 24 25 26
void Distribution::uniform(double lower, double upper, size_t nloc, int seed) {
    
    nloc_m = nloc;
    
27
    std::mt19937_64 mt(0/*seed*/ /*0*/);
frey_m's avatar
frey_m committed
28 29 30
    
    std::uniform_real_distribution<> dist(lower, upper);
    
31 32 33 34 35 36 37 38
//     // assume that seed == rank of node
//     mt.discard(6 * (nloc + 1) * seed);
    
    // assume that seed == rank of node    
    // inefficient but only way to make sure that parallel distribution is equal to sequential
    for (size_t i = 0; i < 6 * nloc_m * seed; ++i)
        dist(mt);
    
frey_m's avatar
frey_m committed
39 40 41 42 43 44 45 46
    x_m.resize(nloc);
    y_m.resize(nloc);
    z_m.resize(nloc);
    
    px_m.resize(nloc);
    py_m.resize(nloc);
    pz_m.resize(nloc);
    
frey_m's avatar
frey_m committed
47 48
    q_m.resize(nloc);
    
frey_m's avatar
frey_m committed
49 50 51 52 53 54 55 56
    for (size_t i = 0; i < nloc; ++i) {
        x_m[i] = dist(mt);
        y_m[i] = dist(mt);
        z_m[i] = dist(mt);
        
        px_m[i] = dist(mt);
        py_m[i] = dist(mt);
        pz_m[i] = dist(mt);
frey_m's avatar
frey_m committed
57 58
        
        q_m[i] = 1.0;
frey_m's avatar
frey_m committed
59 60 61 62 63 64 65
    }
}

void Distribution::gaussian(double mean, double stddev, size_t nloc, int seed) {
    
    nloc_m = nloc;
    
66
    std::mt19937_64 mt(0/*seed*/ /*0*/);
frey_m's avatar
frey_m committed
67 68 69
    
    std::normal_distribution<double> dist(mean, stddev);
    
70 71 72 73 74 75 76 77
//     // assume that seed == rank of node
//     mt.discard(6 * (nloc + 1) * seed);

    // assume that seed == rank of node    
    // inefficient but only way to make sure that parallel distribution is equal to sequential
    for (size_t i = 0; i < 6 * nloc_m * seed; ++i)
        dist(mt);
    
frey_m's avatar
frey_m committed
78 79 80 81 82 83 84 85
    x_m.resize(nloc);
    y_m.resize(nloc);
    z_m.resize(nloc);
    
    px_m.resize(nloc);
    py_m.resize(nloc);
    pz_m.resize(nloc);
    
frey_m's avatar
frey_m committed
86 87
    q_m.resize(nloc);
    
frey_m's avatar
frey_m committed
88 89 90 91 92 93 94 95
    for (size_t i = 0; i < nloc; ++i) {
        x_m[i] = dist(mt);
        y_m[i] = dist(mt);
        z_m[i] = dist(mt);
        
        px_m[i] = dist(mt);
        py_m[i] = dist(mt);
        pz_m[i] = dist(mt);
frey_m's avatar
frey_m committed
96 97 98
        
        
        q_m[i] = 1.0;
frey_m's avatar
frey_m committed
99 100 101 102 103 104 105 106 107 108
    }
}

void Distribution::readH5(const std::string& filename, int step) {
    H5Reader h5(filename);
    
    h5.open(step);
    
    
    
frey_m's avatar
frey_m committed
109
    nloc_m = h5.getNumParticles();
110 111
    ntot_m = nloc_m;
    
frey_m's avatar
frey_m committed
112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128
    size_t numParticlesPerNode = nloc_m / Ippl::getNodes();

    size_t firstParticle = numParticlesPerNode * Ippl::myNode();
    size_t lastParticle = firstParticle + numParticlesPerNode - 1;
    if (Ippl::myNode() == Ippl::getNodes() - 1)
        lastParticle = nloc_m - 1;

    nloc_m = lastParticle - firstParticle + 1;
    
    x_m.resize(nloc_m);
    y_m.resize(nloc_m);
    z_m.resize(nloc_m);
    
    px_m.resize(nloc_m);
    py_m.resize(nloc_m);
    pz_m.resize(nloc_m);
    
frey_m's avatar
frey_m committed
129 130
    q_m.resize(nloc_m);
    
frey_m's avatar
frey_m committed
131 132 133
    h5.read(x_m, px_m,
            y_m, py_m,
            z_m, pz_m,
frey_m's avatar
frey_m committed
134
            q_m,
frey_m's avatar
frey_m committed
135 136 137 138 139 140 141
            firstParticle,
            lastParticle);
    
    h5.close();
}


142
void Distribution::injectBeam(PartBunchBase& bunch, bool doDelete, std::array<double, 3> shift) {
frey_m's avatar
frey_m committed
143
    
144
    // destroy all partcles
145
    if ( doDelete && bunch.getLocalNum() )
146 147
        bunch.destroyAll();
    
148 149 150
    // previous number of particles
    int prevnum = bunch.getLocalNum();

frey_m's avatar
frey_m committed
151 152
    // create memory space
    bunch.create(nloc_m);
153 154

    for (int i = nloc_m - 1; i >= 0; --i) {
155
        bunch.setR(Vector_t(x_m[i] + shift[0], y_m[i] + shift[1], z_m[i] + shift[2]), i + prevnum);
156 157
        bunch.setP(Vector_t(px_m[i], py_m[i], pz_m[i]), i + prevnum);
        bunch.setQM(q_m[i], i + prevnum);
frey_m's avatar
frey_m committed
158 159 160 161 162 163 164
        
        x_m.pop_back();
        y_m.pop_back();
        z_m.pop_back();
        px_m.pop_back();
        py_m.pop_back();
        pz_m.pop_back();
frey_m's avatar
frey_m committed
165 166
        
        q_m.pop_back();
frey_m's avatar
frey_m committed
167
    }
168 169 170 171 172 173 174 175 176 177
}


void Distribution::setDistribution(PartBunchBase& bunch, const std::string& filename, int step) {
    
    readH5(filename, step);
    
    for (unsigned int i = 0; i < bunch.getLocalNum(); ++i)
        bunch.setR(Vector_t(x_m[i], y_m[i], z_m[i]), i);
}
178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206

void Distribution::print2file(std::string pathname) {
    
    for (int n = 0; n < Ippl::getNodes(); ++n) {
        
        if ( n == Ippl::myNode() ) {
            
            std::ofstream out;
            switch (n) {
                case 0:
                    out.open(pathname);
                    out << ntot_m << std::endl;
                    break;
                default:
                    out.open(pathname, std::ios::app);
                    break;
            }
            
            for (std::size_t i = 0; i < x_m.size(); ++i)
                out << x_m[i] << " " << px_m[i] << " "
                    << y_m[i] << " " << py_m[i] << " "
                    << z_m[i] << " " << pz_m[i] << std::endl;
            
            out.close();
        }
        
        Ippl::Comm->barrier();
    }
}