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_;

        uint wid_; 
        uint nlocal_; 

        std::string ID_; 

        Matrix3 H_; 
        Matrix3 Hinv_; 

        Matrix3 virial_; 

        Vector3 origin_; 

        Bool3 isperiodic_; 

        std::vector<Vector3> positions_; 
        std::vector<Integer3> images_; 
        std::vector<Vector3> velocities_; 
        std::vector<Vector3> forces_; 
        std::vector<double> masses_; 
        std::vector<double> charges_; 
        Label atomids_; 
        Label types_; 
        std::vector<std::vector<double> > sigma_; 

        int iteration_; 
        int targetiter_; 
        double temperature_; 
        double energy_; 
        double kb_; 

        double dielectric_; 
        double qqrd2e_; 
        bool changed_; 

    public:

        Snapshot(const MPI_Comm& comm, uint wid) :
        comm_(comm), wid_(wid), H_(), Hinv_(), virial_(Matrix3::Zero()), 
        origin_({0,0,0}), isperiodic_({true, true, true}), positions_(0), 
        images_(0), velocities_(0), forces_(0), masses_(0), atomids_(0), 
        types_(0), iteration_(0), temperature_(0), energy_(0), kb_(0)
        {}


        int GetIteration() const {return iteration_; }
        

        int 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(int 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<Integer3>& GetImageFlags() const { return images_; }

        std::vector<Integer3>& GetImageFlags() 
        { 
            changed_ = true;
            return images_; 
        }


        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_);
        }


        Vector3 UnwrapVector(const Vector3& v, const Integer3& image) const
        {
            return H_*image.cast<double>()+v;
        }


        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) const
        {
            // Get total mass.
            auto mtot = TotalMass(indices);

            return CenterOfMass(indices, mtot);         
        }


        Vector3 CenterOfMass(const Label& indices, double mtot) const
        {
            // Store coorinates 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);

            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(), 0);

            // All-gather data.
            MPI_Allgatherv(pos.data(), pos.size(), MPI_DOUBLE, gpos.data(), pcounts.data(), pdispls.data(), MPI_DOUBLE, comm_);
            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::vector<std::vector<double>>& GetSigmas() const { return sigma_; }

        std::vector<std::vector<double>>& GetSigmas() 
        {
            changed_ = true; 
            return sigma_;
        }


        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(){}
    };
}