Snapshot.h

 
#pragma once
#include <numeric>
#include <vector>
#include <memory>
#include <mxx/comm.hpp>
#include "types.h"
 
namespace SSAGES
{
 
    inline double roundf(double x)
    {
        return ( x >= 0 ) ? floor( x + 0.5 ) : ceil( x - 0.5 );
    }
 
 
    class Snapshot
    {
    private:
        mxx::comm comm_;
 
        unsigned int wid_; 
        unsigned int nlocal_; 
 
        std::string ID_; 
 
        Matrix3 H_; 
        Matrix3 Hinv_; 
 
        Matrix3 virial_; 
 
        Vector3 origin_; 
 
        Bool3 isperiodic_; 
 
        std::vector<Vector3> positions_; 
        std::vector<Vector3> velocities_; 
        std::vector<Vector3> forces_; 
        std::vector<double> masses_; 
        std::vector<double> charges_; 
        Label atomids_; 
        Label types_; 
 
        size_t iteration_; 
        size_t targetiter_; 
        double temperature_; 
        double energy_; 
        double kb_; 
 
        double dielectric_; 
        double qqrd2e_; 
        bool changed_; 
 
    public:
 
        Snapshot(const MPI_Comm& comm, unsigned int wid) :
        comm_(comm), wid_(wid), H_(), Hinv_(), virial_(Matrix3::Zero()), 
        origin_({0,0,0}), isperiodic_({true, true, true}), positions_(0), 
        velocities_(0), forces_(0), masses_(0), atomids_(0), types_(0),
        iteration_(0), temperature_(0), energy_(0), kb_(0)
        {}
 
 
        size_t GetIteration() const {return iteration_; }
        
 
        size_t GetTargetIterations() const {return targetiter_; }
 
 
        double GetTemperature() const {return temperature_; }
 
 
        double GetEnergy() const { return energy_; }
 
 
        const Matrix3& GetHMatrix() const { return H_; }
 
 
        const Matrix3& GetVirial() const { return virial_; }
 
 
        Matrix3& GetVirial() { return virial_; }
 
 
        const Vector3& GetOrigin() const { return origin_; }
 
 
        const Bool3& IsPeriodic() const {return isperiodic_; }
 
 
        double GetVolume() const { return H_.determinant(); }
 
 
        double GetKb() const { return kb_; }
 
 
        double GetDielectric() const { return dielectric_; }
 
 
        double Getqqrd2e() const { return qqrd2e_; }
        
 
        const mxx::comm& GetCommunicator() const
        {
            return comm_;
        }
 
 
        unsigned GetWalkerID() const { return wid_; }
 
 
 
        unsigned GetNumAtoms() const { return nlocal_; }
 
 
        void SetIteration(size_t iteration) 
        {
            iteration_ = iteration; 
            changed_ = true; 
        }
        
 
        void SetTargetIterations(int target) 
        {
            targetiter_ = target; 
            changed_ = true; 
        }
 
 
        void SetTemperature(double temperature) 
        { 
            temperature_ = temperature;
            changed_ = true;
        }
 
 
        void SetEnergy(double energy) 
        {
            energy_ = energy;
            changed_ = true;
        }
 
 
        void SetHMatrix(const Matrix3& hmat)
        {
            H_ = hmat;
            Hinv_ = hmat.inverse();
            changed_ = true;
        }
 
 
        void SetVirial(const Matrix3& virial)
        {
            virial_ = virial; 
            changed_ = true;
        }
 
 
        void SetOrigin(const Vector3& origin)
        {
            origin_ = origin; 
            changed_ = true;
        }
 
 
        void SetPeriodicity(const Bool3& isperiodic)
        {
            isperiodic_ = isperiodic;
            changed_ = true;
        }
 
 
        void SetKb(double kb) 
        {
            kb_ = kb;
            changed_ = true;
        }
 
 
        void SetDielectric(double dielectric) 
        { 
            dielectric_ = dielectric;
            changed_ = true;
        }
 
 
        void Setqqrd2e(double qqrd2e) 
        { 
            qqrd2e_ = qqrd2e;
            changed_ = true;
        }
 
 
        void SetNumAtoms(unsigned int natoms) { nlocal_ = natoms; }
 
 
        const std::vector<Vector3>& GetPositions() const { return positions_; }
 
        std::vector<Vector3>& GetPositions() 
        { 
            changed_ = true;
            return positions_; 
        }
 
 
        const std::vector<Vector3>& GetVelocities() const { return velocities_; }
 
        std::vector<Vector3>& GetVelocities() 
        {
            changed_ = true;
            return velocities_; 
        }
 
 
        const std::vector<Vector3>& GetForces() const { return forces_; }
 
        std::vector<Vector3>& GetForces() 
        {
            changed_ = true; 
            return forces_; 
        }
 
 
        const std::vector<double>& GetMasses() const { return masses_; }
 
        std::vector<double>& GetMasses() 
        {
            changed_ = true; 
            return masses_; 
        }
 
 
        Vector3 ScaleVector(const Vector3& v) const
        {
            return Hinv_*(v-origin_);
        }
 
 
        void ApplyMinimumImage(Vector3* v) const
        {
            Vector3 scaled = Hinv_*(*v);
 
            for(int i = 0; i < 3; ++i)
                scaled[i] -= isperiodic_[i]*roundf(scaled[i]);
 
            *v = H_*scaled;
        }
        
 
        Vector3 ApplyMinimumImage(const Vector3& v) const
        {
            Vector3 scaled = Hinv_*v;
 
            for(int i = 0; i < 3; ++i)
                scaled[i] -= isperiodic_[i]*roundf(scaled[i]);
 
            return H_*scaled;   
        }
 
 
        double TotalMass(const Label& indices) const
        {
            auto mtot = 0.;
            for(auto& i : indices)
                mtot += masses_[i];
            MPI_Allreduce(MPI_IN_PLACE, &mtot, 1, MPI_DOUBLE, MPI_SUM, comm_);
            return mtot;
        }
 
 
        Vector3 CenterOfMass(const Label& indices, bool mass_weight = true) const
        {
            // Store coordinates and masses in vectors to gather.
            std::vector<double> pos, mass, gpos, gmass;
            std::vector<int> pcounts(comm_.size(), 0), mcounts(comm_.size(), 0); 
            std::vector<int> pdispls(comm_.size()+1, 0), mdispls(comm_.size()+1, 0);
 
            // Calculate total mass
            double mtot = mass_weight ? TotalMass(indices) : indices.size();
 
            pcounts[comm_.rank()] = 3*indices.size();
            mcounts[comm_.rank()] = indices.size();
 
            // Reduce counts.
            MPI_Allreduce(MPI_IN_PLACE, pcounts.data(), pcounts.size(), MPI_INT, MPI_SUM, comm_);
            MPI_Allreduce(MPI_IN_PLACE, mcounts.data(), mcounts.size(), MPI_INT, MPI_SUM, comm_);
 
            // Compute displacements.
            std::partial_sum(pcounts.begin(), pcounts.end(), pdispls.begin() + 1);
            std::partial_sum(mcounts.begin(), mcounts.end(), mdispls.begin() + 1);
            
            // Fill up mass and position vectors.
            for(auto& idx : indices)
            {
                auto& p = positions_[idx];
                pos.push_back(p[0]);
                pos.push_back(p[1]);
                pos.push_back(p[2]);
                mass.push_back(masses_[idx]);
            }
 
            // Re-size receiving vectors. 
            gpos.resize(pdispls.back(), 0);
            gmass.resize(mdispls.back(), 1.0); // fictitious mass of 1
 
            // All-gather data.
            MPI_Allgatherv(pos.data(), pos.size(), MPI_DOUBLE, gpos.data(), pcounts.data(), pdispls.data(), MPI_DOUBLE, comm_);
            if(mass_weight) { // overwrite gmass with true mass values
                MPI_Allgatherv(mass.data(), mass.size(), MPI_DOUBLE, gmass.data(), mcounts.data(), mdispls.data(), MPI_DOUBLE, comm_);
            }
 
            // Loop through atoms and compute mass weighted sum.
            // We march linearly through list and find nearest image
            // to each successive particle to properly unwrap object.
            Vector3 ppos = {gpos[0], gpos[1], gpos[2]}; // Previous unwrapped position.
            Vector3 cpos = ppos;
            Vector3 xcm = gmass[0]*cpos;
 
            for(size_t i = 1, j = 3; i < gmass.size(); ++i, j += 3)
            {
                cpos = {gpos[j], gpos[j+1], gpos[j+2]};
                cpos = ApplyMinimumImage(cpos - ppos) + ppos;
                xcm += gmass[i]*cpos;
                ppos = cpos;
            }
 
            return xcm/mtot;
        }
 
 
        const Label& GetAtomIDs() const { return atomids_; }
 
        Label& GetAtomIDs()
        {
            changed_ = true;
            return atomids_;
        }
 
 
        int GetLocalIndex(int id) const
        {
 
            auto s = std::find(atomids_.begin(), atomids_.end(), id);
            if(s == atomids_.end())
                return -1;
            else
                return s - atomids_.begin();
        }
 
 
        void GetLocalIndices(const Label& ids, Label* indices) const
        {
            for(auto& id : ids)
            {
                auto idx = GetLocalIndex(id);
                if(idx != -1)
                    indices->push_back(idx);
            }
        }
 
 
        const std::vector<double>& GetCharges() const { return charges_; }
 
        std::vector<double>& GetCharges()
        {
            changed_ = true;
            return charges_;
        }
        
 
        const Label& GetAtomTypes() const { return types_; }
 
        Label& GetAtomTypes() 
        {
            changed_ = true; 
            return types_; 
        }
 
 
        const std::string& GetSnapshotID() const { return ID_; }
 
        std::string& GetSnapshotID() 
        {
            changed_ = true; 
            return ID_; 
        }
 
 
        bool HasChanged() const { return changed_; }
 
 
        void Changed(bool state) { changed_ = state; }
 
 
        std::vector<double> SerializePositions()
        {
 
            std::vector<int> pcounts(comm_.size(), 0); 
            std::vector<int> pdispls(comm_.size()+1, 0);
 
            pcounts[comm_.rank()] = 3*nlocal_;
 
            // Reduce counts.
            MPI_Allreduce(MPI_IN_PLACE, pcounts.data(), pcounts.size(), MPI_INT, MPI_SUM, comm_);
 
            // Compute displacements.
            std::partial_sum(pcounts.begin(), pcounts.end(), pdispls.begin() + 1);
 
            // Re-size receiving vectors.
            std::vector<double> positions;
            positions.resize(pdispls.back(), 0);
 
            std::vector<double> ptemp;
            
            for(auto& p : positions_)
            {
                ptemp.push_back(p[0]);
                ptemp.push_back(p[1]);
                ptemp.push_back(p[2]);
            }
 
            // All-gather data.
            MPI_Allgatherv(ptemp.data(), ptemp.size(), MPI_DOUBLE, positions.data(), pcounts.data(), pdispls.data(), MPI_DOUBLE, comm_);
            return positions;
        }
 
 
        std::vector<double> SerializeVelocities()
        {
            std::vector<int> vcounts(comm_.size(), 0); 
            std::vector<int> vdispls(comm_.size()+1, 0);
 
            vcounts[comm_.rank()] = 3*nlocal_;
 
            // Reduce counts.
            MPI_Allreduce(MPI_IN_PLACE, vcounts.data(), vcounts.size(), MPI_INT, MPI_SUM, comm_);
 
            // Compute displacements.
            std::partial_sum(vcounts.begin(), vcounts.end(), vdispls.begin() + 1);
 
            // Re-size receiving vectors.
            std::vector<double> velocities;
            velocities.resize(vdispls.back(), 0);
 
            std::vector<double> vtemp;
            
            for(auto& v : velocities_)
            {
                vtemp.push_back(v[0]);
                vtemp.push_back(v[1]);
                vtemp.push_back(v[2]);
            }
 
            // All-gather data.
            MPI_Allgatherv(vtemp.data(), vtemp.size(), MPI_DOUBLE, velocities.data(), vcounts.data(), vdispls.data(), MPI_DOUBLE, comm_);
            return velocities;
        }
 
 
        std::vector<int> SerializeIDs()
        {
            std::vector<int> mcounts(comm_.size(), 0); 
            std::vector<int> mdispls(comm_.size()+1, 0);
 
            mcounts[comm_.rank()] = nlocal_;
 
            // Reduce counts.
            MPI_Allreduce(MPI_IN_PLACE, mcounts.data(), mcounts.size(), MPI_INT, MPI_SUM, comm_);
 
            // Compute displacements.
            std::partial_sum(mcounts.begin(), mcounts.end(), mdispls.begin() + 1);
 
            // Re-size receiving vectors.
            std::vector<int> IDs;
            IDs.resize(mdispls.back(), 0);
 
            // All-gather data.
            MPI_Allgatherv(atomids_.data(), atomids_.size(), MPI_INT, IDs.data(), mcounts.data(), mdispls.data(), MPI_INT, comm_);
            return IDs;
        }
 
        ~Snapshot(){}
    };
}