initial version for moments computer

src/Algorithms/ParallelTTracker.cpp

@@ -469,15 +469,14 @@ void ParallelTTracker::emitParticles(long long step) {
 
 
 void ParallelTTracker::computeSpaceChargeFields(unsigned long long step) {
-    if (numParticlesInSimulation_m <= minBinEmitted_m) return;
-
-    if (!itsBunch_m->hasFieldSolver()) return;
+    if (numParticlesInSimulation_m <= minBinEmitted_m || !itsBunch_m->hasFieldSolver()) {
+        return;
+    }
 
     itsBunch_m->calcBeamParameters();
     Quaternion alignment = getQuaternion(itsBunch_m->get_pmean(), Vector_t(0, 0, 1));
     CoordinateSystemTrafo beamToReferenceCSTrafo(Vector_t(0, 0, pathLength_m), alignment.conjugate());
     CoordinateSystemTrafo referenceToBeamCSTrafo = beamToReferenceCSTrafo.inverted();
-
     const unsigned int localNum1 = itsBunch_m->getLocalNum();
     for (unsigned int i = 0; i < localNum1; ++ i) {
         itsBunch_m->R[i] = referenceToBeamCSTrafo.transformTo(itsBunch_m->R[i]);
@@ -1377,4 +1376,4 @@ void ParallelTTracker::evenlyDistributeParticles() {
     if (requests.size() > 0) {
         MPI_Waitall(requests.size(), &(requests[0]), MPI_STATUSES_IGNORE);
     }
-}
+}
\ No newline at end of file

src/Classic/Algorithms/DistributionMoments.cpp

@@ -4,21 +4,73 @@
 #include "Utilities/Util.h"
 
 #include "Message/GlobalComm.h"
+#include "Utility/Inform.h"
+
+#include "OpalParticle.h"
+#include "PartBunchBase.h"
+
+extern Inform* gmsg;
 
 DistributionMoments::DistributionMoments():
     meanR_m(0.0),
     meanP_m(0.0),
     stdR_m(0.0),
     stdP_m(0.0),
+    stdRP_m(0.0),
     normalizedEps_m(0.0),
     geometricEps_m(0.0),
     meanKineticEnergy_m(0.0),
+    stdKineticEnergy_m(0.0),
+    Dx_m(0.0),
+    DDx_m(0.0),
+    Dy_m(0.0),
+    DDy_m(0.0),
     moments_m(),
     totalNumParticles_m(0)
 {
     std::fill(std::begin(centroid_m), std::end(centroid_m), 0.0);
 }
 
+void DistributionMoments::compute(PartBunchBase<double, 3> const& bunch)
+{
+    reset();
+
+    totalNumParticles_m = bunch.getTotalNum();
+
+    if (totalNumParticles_m == 0) {
+        return;
+    }
+
+    computeMoments(bunch);
+
+    Vector_t squaredEps, fac, squaredSumR, squaredSumP, sumRP;
+    double perParticle = 1.0 / totalNumParticles_m;
+    for (unsigned int i = 0; i < 3; ++ i) {
+        meanR_m(i) = centroid_m[2 * i] * perParticle;
+        meanP_m(i) = centroid_m[2 * i + 1] * perParticle;
+        squaredSumR(i) = moments_m(2 * i, 2 * i) - totalNumParticles_m * std::pow(meanR_m(i), 2);
+        squaredSumP(i) = std::max(0.0,
+                                  moments_m(2 * i + 1, 2 * i + 1) - totalNumParticles_m * std::pow(meanP_m(i), 2));
+        sumRP(i) = moments_m(2 * i, 2 * i + 1) - totalNumParticles_m * meanR_m(i) * meanP_m(i);
+    }
+
+    squaredEps = (squaredSumR * squaredSumP - sumRP * sumRP) * std::pow(perParticle, 2);
+    sumRP *= perParticle;
+
+    for (unsigned int i = 0; i < 3; ++ i) {
+        stdR_m(i) = std::sqrt(squaredSumR(i) * perParticle);
+        stdP_m(i) = std::sqrt(squaredSumP(i) * perParticle);
+        normalizedEps_m(i) = std::sqrt(std::max(squaredEps(i), 0.0));
+        double tmp = stdR_m(i) * stdP_m(i);
+        fac(i) = (std::abs(tmp) < 1e-10) ? 0.0: 1.0 / tmp;
+    }
+
+    stdRP_m = sumRP * fac;
+
+    double betaGamma = std::sqrt(std::pow(meanGamma_m, 2) - 1.0);
+    geometricEps_m = normalizedEps_m / Vector_t(betaGamma);
+}
+
 void DistributionMoments::compute(std::vector<OpalParticle> const& particles)
 {
     reset();
@@ -44,7 +96,7 @@ void DistributionMoments::compute(std::vector<OpalParticle> const& particles)
     }
 
     squaredEps = (squaredSumR * squaredSumP - sumRP * sumRP) * std::pow(perParticle, 2);
-    // sumRP *= perParticle;
+    sumRP *= perParticle;
 
     for (unsigned int i = 0; i < 3; ++ i) {
         stdR_m(i) = std::sqrt(squaredSumR(i) * perParticle);
@@ -54,15 +106,28 @@ void DistributionMoments::compute(std::vector<OpalParticle> const& particles)
         fac(i) = (std::abs(tmp) < 1e-10) ? 0.0: 1.0 / tmp;
     }
 
-    // stdRP_m = sumRP * fac;
+    stdRP_m = sumRP * fac;
 
     double betaGamma = std::sqrt(std::pow(meanGamma_m, 2) - 1.0);
     geometricEps_m = normalizedEps_m / Vector_t(betaGamma);
 }
 
-void DistributionMoments::computeMoments(std::vector<OpalParticle> const& particles)
+void DistributionMoments::computeMeanKineticEnergy(PartBunchBase<double, 3> const& bunch)
+{
+    double data[] = {0.0, 0.0};
+    for (OpalParticle const& particle: bunch) {
+        data[0] += Util::getKineticEnergy(particle.P(), particle.mass());
+    }
+    data[1] = bunch.getLocalNum();
+    allreduce(data, 2, std::plus<double>());
+
+    meanKineticEnergy_m = data[0] / data[1];
+}
+
+template<class Container>
+void DistributionMoments::computeMoments(Container const& particles)
 {
-    std::vector<double> localMoments(35);
+    std::vector<double> localMoments(36);
 
     for (OpalParticle const& particle: particles) {
         unsigned int l = 6;
@@ -79,11 +144,12 @@ void DistributionMoments::computeMoments(std::vector<OpalParticle> const& partic
             localMoments[l + 3] += r2 * r2;
         }
         double gamma = Util::getGamma(particle.P());
-        localMoments[33] = (gamma - 1.0) * particle.mass();
-        localMoments[34] = gamma;
+        double eKin = (gamma - 1.0) * particle.mass();
+        localMoments[33] += eKin;
+        localMoments[34] += std::pow(eKin, 2);
+        localMoments[35] += gamma;
     }
 
-
     allreduce(localMoments.data(), localMoments.size(), std::plus<double>());
 
     for (unsigned int i = 0; i < 6; ++ i) {
@@ -112,7 +178,8 @@ void DistributionMoments::computeMoments(std::vector<OpalParticle> const& partic
     }
 
     meanKineticEnergy_m = localMoments[33] * perParticle;
-    meanGamma_m = localMoments[34] * perParticle;
+    stdKineticEnergy_m = localMoments[34] * perParticle;
+    meanGamma_m = localMoments[35] * perParticle;
 }
 
 void DistributionMoments::reset()
@@ -121,11 +188,17 @@ void DistributionMoments::reset()
     meanP_m = 0.0;
     stdR_m = 0.0;
     stdP_m = 0.0;
+    stdRP_m = 0.0;
     normalizedEps_m = 0.0;
     geometricEps_m = 0.0;
     halo_m = 0.0;
 
     meanKineticEnergy_m = 0.0;
+    stdKineticEnergy_m = 0.0;
+    Dx_m = 0.0;
+    DDx_m = 0.0;
+    Dy_m = 0.0;
+    DDy_m = 0.0;
     std::fill(std::begin(centroid_m), std::end(centroid_m), 0.0);
     moments_m = FMatrix<double, 6, 6>(0.0);
 

src/Classic/Algorithms/DistributionMoments.h

 #ifndef DISTRIBUTIONMOMENTS_H
 #define DISTRIBUTIONMOMENTS_H
 
-#include "OpalParticle.h"
-
 #include "FixedAlgebra/FMatrix.h"
 
+#include "Vektor.h"
+
 #include <vector>
 
+class OpalParticle;
+template<class T, unsigned Dim>
+class PartBunchBase;
+
 class DistributionMoments {
 public:
     DistributionMoments();
 
     void compute(std::vector<OpalParticle> const&);
+    void compute(PartBunchBase<double, 3> const&);
+    void computeMeanKineticEnergy(PartBunchBase<double, 3> const&);
 
     Vector_t getMeanPosition() const;
     Vector_t getStandardDeviationPosition() const;
@@ -19,22 +25,38 @@ public:
     Vector_t getStandardDeviationMomentum() const;
     Vector_t getNormalizedEmittance() const;
     Vector_t getGeometricEmittance() const;
+    Vector_t getStandardDeviationRP() const;
+    Vector_t getHalo() const;
+    double getMeanGamma() const;
     double getMeanKineticEnergy() const;
+    double getStandardDeviationKineticEnergy() const;
+    double getDx() const;
+    double getDDx() const;
+    double getDy() const;
+    double getDDy() const;
+
 
 private:
-    void computeMoments(std::vector<OpalParticle> const&);
+    template<class Container>
+    void computeMoments(Container const&);
     void reset();
 
     Vector_t meanR_m;
     Vector_t meanP_m;
     Vector_t stdR_m;
     Vector_t stdP_m;
+    Vector_t stdRP_m;
     Vector_t normalizedEps_m;
     Vector_t geometricEps_m;
     Vector_t halo_m;
 
     double meanKineticEnergy_m;
+    double stdKineticEnergy_m;
     double meanGamma_m;
+    double Dx_m;
+    double DDx_m;
+    double Dy_m;
+    double DDy_m;
     double centroid_m[6];
     FMatrix<double, 6, 6> moments_m;
 
@@ -77,10 +99,58 @@ Vector_t DistributionMoments::getGeometricEmittance() const
     return geometricEps_m;
 }
 
+inline
+Vector_t DistributionMoments::getStandardDeviationRP() const
+{
+    return stdRP_m;
+}
+
+inline
+Vector_t DistributionMoments::getHalo() const
+{
+    return halo_m;
+}
+
+inline
+double DistributionMoments::getMeanGamma() const
+{
+    return meanGamma_m;
+}
+
 inline
 double DistributionMoments::getMeanKineticEnergy() const
 {
     return meanKineticEnergy_m;
 }
 
+inline
+double DistributionMoments::getStandardDeviationKineticEnergy() const
+{
+    return stdKineticEnergy_m;
+}
+
+inline
+double DistributionMoments::getDx() const
+{
+    return Dx_m;
+}
+
+inline
+double DistributionMoments::getDDx() const
+{
+    return DDx_m;
+}
+
+inline
+double DistributionMoments::getDy() const
+{
+    return Dy_m;
+}
+
+inline
+double DistributionMoments::getDDy() const
+{
+    return DDy_m;
+}
+
 #endif
\ No newline at end of file

src/Classic/Algorithms/OpalParticle.h

@@ -25,7 +25,7 @@ class OpalParticle {
 public:
 
     // Particle coordinate numbers.
-    enum { X, PX, Y, PY, T, PT };
+    enum { X, Y, T, PX, PY, PT };
 
     /// Constructor.
     //  Construct particle with the given coordinates.

src/Classic/Algorithms/PartBunchBase.h

@@ -18,10 +18,12 @@
 #ifndef PART_BUNCH_BASE_H
 #define PART_BUNCH_BASE_H
 
-#include "Utility/IpplTimings.h"
 #include "Particle/AbstractParticle.h"
 #include "Particle/ParticleAttrib.h"
+#include "Utility/IpplTimings.h"
+#include "Utilities/GeneralClassicException.h"
 
+#include "Algorithms/DistributionMoments.h"
 #include "Algorithms/CoordinateSystemTrafo.h"
 #include "Algorithms/OpalParticle.h"
 #include "Algorithms/PBunchDefs.h"
@@ -50,7 +52,7 @@ namespace ParticleType {
 }
 
 template <class T, unsigned Dim>
-class PartBunchBase
+class PartBunchBase : std::enable_shared_from_this<PartBunchBase<T, Dim>>
 {
 public:
     typedef typename AbstractParticle<T, Dim>::ParticlePos_t   ParticlePos_t;
@@ -215,6 +217,72 @@ public:
 
     OpalParticle get_part(int ii);
 
+    class ConstIterator {
+        friend class PartBunchBase<T, Dim>;
+
+    public:
+        ConstIterator():
+            bunch_m(nullptr),
+            index_m(0)
+        {}
+        ConstIterator(PartBunchBase const* bunch, unsigned int i):
+            bunch_m(bunch),
+            index_m(i)
+        {}
+
+        ~ConstIterator()
+        {
+            bunch_m = nullptr;
+        }
+
+        bool operator == (ConstIterator const& rhs) const
+        {
+            return bunch_m == rhs.bunch_m && index_m == rhs.index_m;
+        }
+
+        bool operator != (ConstIterator const& rhs) const
+        {
+            return bunch_m != rhs.bunch_m || index_m != rhs.index_m;
+        }
+
+        OpalParticle operator*() const
+        {
+            if (index_m >= bunch_m->getLocalNum()) {
+                throw GeneralClassicException("PartBunchBase::ConstIterator::operator*", "out of bouds");
+            }
+
+            return OpalParticle(bunch_m->R[index_m](0), bunch_m->P[index_m](0),
+                                bunch_m->R[index_m](1), bunch_m->P[index_m](1),
+                                bunch_m->R[index_m](2), bunch_m->P[index_m](2),
+                                bunch_m->Q[index_m], bunch_m->M[index_m]);
+        }
+
+        ConstIterator operator++()
+        {
+            ++index_m;
+            return *this;
+        }
+
+        ConstIterator operator++(int)
+        {
+            ConstIterator it = *this;
+            ++index_m;
+
+            return it;
+        }
+    private:
+        PartBunchBase const* bunch_m;
+        unsigned int index_m;
+    };
+
+    ConstIterator begin() const {
+        return ConstIterator(this, 0);
+    }
+
+    ConstIterator end() const {
+        return ConstIterator(this, getLocalNum());
+    }
+
     /// Return maximum amplitudes.
     //  The matrix [b]D[/b] is used to normalise the first two modes.
     //  The maximum normalised amplitudes for these modes are stored
@@ -570,10 +638,10 @@ protected:
     double dt_m;
     /// holds the actual time of the integration
     double t_m;
-    /// mean energy of the bunch (MeV)
-    double eKin_m;
+    // /// mean energy of the bunch (MeV)
+    // double eKin_m;
     /// energy spread of the beam in MeV
-    double dE_m;
+    // double dE_m;
     /// the position along design trajectory
     double spos_m;
 
@@ -590,33 +658,33 @@ protected:
     /// minimal extend of particles
     Vector_t rmin_m;
 
-    /// rms beam size (m)
-    Vector_t rrms_m;
-    /// rms momenta
-    Vector_t prms_m;
-    /// mean position (m)
-    Vector_t rmean_m;
-    /// mean momenta
-    Vector_t pmean_m;
+    // /// rms beam size (m)
+    // Vector_t rrms_m;
+    // /// rms momenta
+    // Vector_t prms_m;
+    // /// mean position (m)
+    // Vector_t rmean_m;
+    // /// mean momenta
+    // Vector_t pmean_m;
 
-    /// rms emittance (not normalized)
-    Vector_t eps_m;
+    // /// rms emittance (not normalized)
+    // Vector_t eps_m;
 
-    /// rms normalized emittance
-    Vector_t eps_norm_m;
+    // /// rms normalized emittance
+    // Vector_t eps_norm_m;
 
-    Vector_t halo_m;
+    // Vector_t halo_m;
 
-    /// rms correlation
-    Vector_t rprms_m;
+    // /// rms correlation
+    // Vector_t rprms_m;
 
-    /// dispersion x & y
-    double Dx_m;
-    double Dy_m;
+    // /// dispersion x & y
+    // double Dx_m;
+    // double Dy_m;
 
-    /// derivative of the dispersion
-    double DDx_m;
-    double DDy_m;
+    // /// derivative of the dispersion
+    // double DDx_m;
+    // double DDy_m;
 
     /// meshspacing of cartesian mesh
     Vector_t hr_m;
@@ -676,6 +744,7 @@ protected:
 
 
     Distribution *dist_m;
+    DistributionMoments momentsComputer_m;
 
     // flag to tell if we are a DC-beam
     bool dcBeam_m;
@@ -683,6 +752,16 @@ protected:
     std::shared_ptr<AbstractParticle<T, Dim> > pbase_m;
 };
 
+template<class T, unsigned Dim>
+typename PartBunchBase<T, Dim>::ConstIterator begin(PartBunchBase<T, Dim> const& bunch) {
+    return bunch.begin();
+}
+
+template<class T, unsigned Dim>
+typename PartBunchBase<T, Dim>::ConstIterator end(PartBunchBase<T, Dim> const& bunch) {
+    return bunch.end();
+}
+
 #include "PartBunchBase.hpp"
 
 #endif
\ No newline at end of file