enforce mean zero by translation, set m/q after particle creation

src/Distribution/Distribution.cpp

@@ -147,10 +147,8 @@ void Distribution::create(size_t& numberOfParticles, double massIneV, double cha
      * 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.
+     * For now we just substract mean to make sure of averages are 0.0
      */
-    size_t numberOfLocalParticles = numberOfParticles;
-    numberOfLocalParticles        = (numberOfParticles + 1) / 2;
 
     size_t mySeed = Options::seed;
 
@@ -165,7 +163,7 @@ void Distribution::create(size_t& numberOfParticles, double massIneV, double cha
 
     switch (distrTypeT_m) {
         case DistributionType::GAUSS:
-            createDistributionGauss(numberOfLocalParticles, massIneV, R, P);
+            createDistributionGauss(numberOfParticles, massIneV, R, P);
             break;
         default:
             throw OpalException("Distribution::create", "Unknown \"TYPE\" of \"DISTRIBUTION\"");
@@ -239,20 +237,91 @@ void Distribution::createDistributionGauss(size_t numberOfParticles, double mass
         mu[i] = 0.0;
         sd[i] = sigmaR_m[i];
     }
-    view_type* Rview = &(R.getView());
+    view_type Rview = R.getView();
     Kokkos::parallel_for(
-            numberOfParticles, ippl::random::randn<double, 3>(*Rview, rand_pool64, mu, sd)
+            numberOfParticles, ippl::random::randn<double, 3>(Rview, rand_pool64, mu, sd)
     );
 
+    double meanR[3], loc_meanR[3];
+    for(int i=0; i<3; i++){
+       meanR[i] = 0.0;
+       loc_meanR[i] = 0.0;
+    }
+    Kokkos::parallel_reduce(
+        "calc moments of particle distr.", numberOfParticles,
+        KOKKOS_LAMBDA(
+                    const int k, double& cent0, double& cent1, double& cent2) {
+                    cent0 += Rview(k)[0];
+                    cent1 += Rview(k)[1];
+                    cent2 += Rview(k)[2];
+        },
+        Kokkos::Sum<double>(loc_meanR[0]), Kokkos::Sum<double>(loc_meanR[1]), Kokkos::Sum<double>(loc_meanR[2]));
+    Kokkos::fence();
+    ippl::Comm->barrier();
+
+    MPI_Allreduce(loc_meanR, meanR, 3, MPI_DOUBLE, MPI_SUM, ippl::Comm->getCommunicator());
+
+    for(int i=0; i<3; i++){
+       meanR[i] = meanR[i]/(1.*numberOfParticles);
+    }
+
+    Kokkos::parallel_for(
+            numberOfParticles,KOKKOS_LAMBDA(
+                    const int k) {
+                    Rview(k)[0] -= meanR[0];
+                    Rview(k)[1] -= meanR[1];
+                    Rview(k)[2] -= meanR[2];
+        }
+    );
+    Kokkos::fence();
+    ippl::Comm->barrier();
+
     // sample P
     for(int i=0; i<3; i++){
         mu[i] = 0.0;
         sd[i] = sigmaP_m[i];
     }
-    view_type* Pview = &(P.getView());
+    view_type Pview = P.getView();
     Kokkos::parallel_for(
-            numberOfParticles, ippl::random::randn<double, 3>(*Pview, rand_pool64, mu, sd)
+            numberOfParticles, ippl::random::randn<double, 3>(Pview, rand_pool64, mu, sd)
     );
+    Kokkos::fence();
+    ippl::Comm->barrier();
+
+    double meanP[3], loc_meanP[3];
+    for(int i=0; i<3; i++){
+       meanP[i] = 0.0;
+       loc_meanP[i] = 0.0;
+    }
+    Kokkos::parallel_reduce(
+        "calc moments of particle distr.", numberOfParticles,
+        KOKKOS_LAMBDA(
+                    const int k, double& cent0, double& cent1, double& cent2) {
+                    cent0 += Pview(k)[0];
+                    cent1 += Pview(k)[1];
+                    cent2 += Pview(k)[2];
+        },
+	Kokkos::Sum<double>(loc_meanP[0]), Kokkos::Sum<double>(loc_meanP[1]), Kokkos::Sum<double>(loc_meanP[2]));
+    Kokkos::fence();
+    ippl::Comm->barrier();
+
+    MPI_Allreduce(loc_meanP, meanP, 3, MPI_DOUBLE, MPI_SUM, ippl::Comm->getCommunicator());
+
+    for(int i=0; i<3; i++){
+       meanP[i] = meanP[i]/(1.*numberOfParticles);
+    }
+
+    Kokkos::parallel_for(
+            numberOfParticles,KOKKOS_LAMBDA(
+                    const int k) {
+                    Pview(k)[0] -= meanP[0];
+                    Pview(k)[1] -= meanP[1];
+                    Pview(k)[2] -= meanP[2];
+        }
+    );
+    Kokkos::fence();
+    ippl::Comm->barrier();
+
 }
 
 void Distribution::printDist(Inform& os, size_t numberOfParticles) const {

src/PartBunch/PartBunch.hpp

@@ -31,11 +31,11 @@ Main loop
      updateExternalFields()
      update()
 
-updateExternalFields(): check if bunch has access to the fields of eternal elements. Maybee phase
+updateExternalFields(): check if bunch has access to the fields of eternal elements. Maybe phase
 out some elements and read in new elements
 
 
-diagnostics(): calculate statistics and maybee write tp h5 and stat files
+diagnostics(): calculate statistics and maybe write tp h5 and stat files
 
 */
 
@@ -691,8 +691,8 @@ public:
 
         std::shared_ptr<ParticleContainer_t> pc = this->pcontainer_m;
 
-        auto Rview = pc->R.getView();
-        auto Pview = pc->P.getView();
+        auto &Rview = pc->R.getView();
+        auto &Pview = pc->P.getView();
 
         ////////////////////////////////////
         //// Calculate Moments of R and P //
@@ -710,9 +710,10 @@ public:
 
         for (unsigned i = 0; i < 2 * Dim; i++) {
             loc_centroid[i] = 0.0;
-            for (unsigned j = 0; j <= i; j++) {
+            centroid[i] = 0.0;
+            for (unsigned j = 0; j < 2 * Dim; j++) {
                 loc_moment[i][j] = 0.0;
-                loc_moment[j][i] = 0.0;
+                moment[i][j] = 0.0;
             }
         }
 
@@ -752,6 +753,13 @@ public:
         MPI_Allreduce(
             loc_centroid, centroid, 2 * Dim, MPI_DOUBLE, MPI_SUM, ippl::Comm->getCommunicator());
 
+        for (unsigned i = 0; i < 2 * Dim; i++) {
+            centroid[i] /= globNp;
+            for (unsigned j = 0; j < 2 * Dim; j++) {
+                moment[i][j] /= globNp;
+            }
+        }
+
         //////////////////////////////////////
         //// Calculate Normalized Emittance //
         //////////////////////////////////////
@@ -937,6 +945,37 @@ public:
             csvout << rrms(0) << "\t" << rrms(1) << "\t" << rrms(2) << "\t"
                    << rmin[0] << "\t" << rmin[1] << "\t" << rmin[2] << "\t" << ippl::Comm->size()
                    << endl;
+            csvout.flush();
+
+            // write the means
+            std::stringstream fname_means;
+            fname_means << "OPAL-X-MEANS";
+            fname_means << ippl::Comm->rank();
+            fname_means << ".csv";
+            Inform csvout_means(NULL, (fname_means.str()).c_str(), Inform::APPEND);
+            csvout_means.precision(6);
+            for(unsigned int i=0; i<2*Dim; i++){
+                csvout_means << centroid[i] << " ";
+            }
+            csvout_means.flush();
+
+            // write the 2nd order moments
+            std::stringstream fname_moments;
+            fname_moments.str(std::string());
+            fname_moments << "OPAL-X-MOMENTS";
+            fname_moments << ippl::Comm->rank();
+            fname_moments << ".csv";
+            Inform csvout_moments(NULL, (fname_moments.str()).c_str(), Inform::APPEND);
+            csvout_moments.precision(6);
+            for(unsigned int i=0; i<2*Dim; i++){
+                for(unsigned int j=0; j<2*Dim; j++){
+                    csvout_moments << moment[i][j] << " ";
+                }
+                csvout_moments << endl;
+            }
+            csvout_moments.flush();
+
+            // write the covariance matrix
 
             /*
                    << "rmeanX,rmeanY,rmeanZ,"
@@ -968,8 +1007,10 @@ public:
     }
 
     void setCharge(double q) {
+         this->pcontainer_m->Q = q;
     }
     void setMass(double mass) {
+         this->pcontainer_m->M = mass;
     }
 
     double getCharge() const {

src/PartBunch/ParticleContainer.hpp

@@ -17,7 +17,7 @@ class ParticleContainer : public ippl::ParticleBase<ippl::ParticleSpatialLayout<
     using Base = ippl::ParticleBase<ippl::ParticleSpatialLayout<T, Dim>>;
 
 public:
-    /// charge in [Cb[
+    /// charge in [Cb]
     ippl::ParticleAttrib<double> Q;
 
     /// mass

src/Track/TrackRun.cpp

@@ -261,8 +261,8 @@ void TrackRun::execute() {
     bunch_m->setT(0.005);
     bunch_m->setBeamFrequency(beam->getFrequency() * Units::MHz2Hz);
     bunch_m->setPType(beam->getParticleName());
-    bunch_m->setCharge(macrocharge_m);
-    bunch_m->setMass(macromass_m);
+    //bunch_m->setCharge(macrocharge_m); //call after creating particles
+    //bunch_m->setMass(macromass_m); //call after creating particles
     //bunch_m->print(*gmsg); // call print of bunch after particle generation
 
     setupBoundaryGeometry();
@@ -299,11 +299,19 @@ void TrackRun::execute() {
           Not sure where to put it. For now, I call it here.
           Alternatively, we can pass pointer to particle container as argument to the opalx's distribution::create, and access R,P,.create() via that
     */
+    // find local number of particles
     size_t nlocal = dist_m->getNumOfLocalParticlesToCreate( beam->getNumberOfParticles() );
+    // allocate memory from IPPL
+    bunch_m->getParticleContainer()->create(nlocal);
+    // set distribution type
     dist_m->setDistType();
-    bunch_m->getParticleContainer()->create(nlocal); // this would allocate memory, etc (from IPPL)
-    dist_m->create(nlocal, 1, 1, bunch_m->getParticleContainer()->R, bunch_m->getParticleContainer()->P); // this would sample particles
+    // sample particles
+    dist_m->create(nlocal, 1, 1, bunch_m->getParticleContainer()->R, bunch_m->getParticleContainer()->P);
+    // set the charge and mass of each particle
+    bunch_m->setCharge(macrocharge_m);
+    bunch_m->setMass(macromass_m);
 
+    dist_m->printInfo(*gmsg);
     bunch_m->print(*gmsg);
 
     // initial statistical data are calculated (rms, eps etc.)