PairwiseCV.h

#pragma once 

#include "CVs/CollectiveVariable.h"
#include "Validator/ObjectRequirement.h"
#include "Drivers/DriverException.h"
#include "Utility/PairwiseKernel.h"
#include "Snapshot.h"
#include "schema.h"
 
namespace SSAGES
{

    class PairwiseCV : public CollectiveVariable
    {
    private:
        Label group1_; 
        Label group2_; 
        PairwiseKernel* pk_; 

    public:

        PairwiseCV(const Label& group1, const Label& group2, PairwiseKernel* pk) : 
        group1_(group1), group2_(group2), pk_(pk)
        {
        }


        void Initialize(const Snapshot& snapshot) override
        {
            using std::to_string;

            auto n1 = group1_.size(), n2 = group2_.size();

            // Make sure atom ID's are on at least one processor. 
            std::vector<int> found1(n1, 0), found2(n2, 0);
            for(size_t i = 0; i < n1; ++i)
            {
                if(snapshot.GetLocalIndex(group1_[i]) != -1)
                    found1[i] = 1;
            }
            
            for(size_t i = 0; i < n2; ++i)
            {
                if(snapshot.GetLocalIndex(group2_[i]) != -1)
                    found2[i] = 1;
            }

            MPI_Allreduce(MPI_IN_PLACE, found1.data(), n1, MPI_INT, MPI_SUM, snapshot.GetCommunicator());
            MPI_Allreduce(MPI_IN_PLACE, found2.data(), n2, MPI_INT, MPI_SUM, snapshot.GetCommunicator());

            unsigned ntot1 = std::accumulate(found1.begin(), found1.end(), 0, std::plus<int>());
            unsigned ntot2 = std::accumulate(found2.begin(), found2.end(), 0, std::plus<int>());
            if(ntot1 != n1)
                throw BuildException({
                    "PairwiseCV: Expected to find " + 
                    to_string(n1) + 
                    " atoms in group 1, but only found " + 
                    to_string(ntot1) + "."
                });         

            if(ntot2 != n2)
                throw BuildException({
                    "PairwiseCV: Expected to find " + 
                    to_string(n2) + 
                    " atoms in group 2, but only found " + 
                    to_string(ntot2) + "."
                });         
        }


        void Evaluate(const Snapshot& snapshot) override
        {
            // Get local atom indices.
            std::vector<int> idx1, idx2;
            snapshot.GetLocalIndices(group1_, &idx1);
            snapshot.GetLocalIndices(group2_, &idx2);

            // Get data from snapshot. 
            auto n = snapshot.GetNumAtoms();
            auto& atomids = snapshot.GetAtomIDs();
            auto& positions = snapshot.GetPositions();
            auto& comm = snapshot.GetCommunicator();

            // Initialize gradient.
            std::fill(grad_.begin(), grad_.end(), Vector3{0,0,0});
            grad_.resize(n, Vector3{0,0,0});
            boxgrad_ = Matrix3::Zero();

            // Fill eigen matrix with coordinate data. 
            std::vector<double> pos1(3*idx1.size()), pos2(3*idx2.size());
            std::vector<int> id1(idx1.size()), id2(idx2.size()); 
            for(size_t i = 0; i < idx1.size(); ++i)
            {
                pos1[3*i+0] = positions[idx1[i]][0]; 
                pos1[3*i+1] = positions[idx1[i]][1]; 
                pos1[3*i+2] = positions[idx1[i]][2];

                id1[i] = atomids[idx1[i]];
            }

            for(size_t i = 0; i < idx2.size(); ++i)
            {
                pos2[3*i+0] = positions[idx2[i]][0]; 
                pos2[3*i+1] = positions[idx2[i]][1]; 
                pos2[3*i+2] = positions[idx2[i]][2];

                id2[i] = atomids[idx2[i]];
            }

            // Gather across all procs. 
            pos1 = mxx::allgatherv(pos1.data(), pos1.size(), comm);
            pos2 = mxx::allgatherv(pos2.data(), pos2.size(), comm);
            id1 = mxx::allgatherv(id1.data(), id1.size(), comm);
            id2 = mxx::allgatherv(id2.data(), id2.size(), comm);

            // Create Eigen map for ease of use.
            using Map = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, 3, Eigen::RowMajor>>;
            Map gpos1(pos1.data(), group1_.size(), 3), gpos2(pos2.data(), group2_.size(), 3);

            val_ = 0;
            double df = 0.;
            // Compute gradient and value.
            for(size_t i = 0; i < group1_.size(); ++i)
            {
                auto t1 = id1[i];
                const auto& pi = gpos1.row(i);
                for(size_t j = 0; j < group2_.size(); ++j)
                {
                    auto t2 = id2[j];
                    const auto& pj = gpos2.row(j);

                    // Skip identical pairs. 
                    if(t1 == t2)
                        continue;
                    
                    // Compute distance and pairwise function. 
                    Vector3 rij = pi - pj;
                    snapshot.ApplyMinimumImage(&rij);
                    auto r = rij.norm();
                    val_ +=  pk_->Evaluate(r, df);

                    // Get local indices and sum gradients.
                    auto lid1 = snapshot.GetLocalIndex(t1);
                    if(lid1 != -1)
                        grad_[lid1] += df*rij/r;
                    
                    auto lid2 = snapshot.GetLocalIndex(t2);
                    if(lid2 != -1)
                        grad_[lid2] -= df*rij/r;

                    // Only sum box gradient on a single proc. 
                    if(comm.rank() == 0)
                        boxgrad_ += df*rij/r*rij.transpose();
                }
            }
        }

        static PairwiseCV* Build(const Json::Value& json, const std::string& path)
        {
            Json::ObjectRequirement validator;
            Json::Value schema;
            Json::Reader reader;

            reader.parse(JsonSchema::PairwiseCV, schema);
            validator.Parse(schema, path);

            // Validate inputs.
            validator.Validate(json, path);
            if(validator.HasErrors())
                throw BuildException(validator.GetErrors());
            
            std::vector<int> group1, group2;
            
            for(auto& s : json["group1"])
                group1.push_back(s.asInt());

            for(auto& s : json["group2"])
                group2.push_back(s.asInt());
            
            return new PairwiseCV(group1, group2, PairwiseKernel::Build(json["kernel"], path));
        }

        ~PairwiseCV()
        {
            delete pk_;
            /*
             * deleting object of abstract class type ‘SSAGES::PairwiseKernel’
             * which has non-virtual destructor will cause undefined behaviour
             * [-Wdelete-non-virtual-dtor]
             *
             */
        }
    };
}