ParticleCoordinateCV.h

#pragma once 

#include "CollectiveVariable.h"
#include "Validator/ObjectRequirement.h"
#include "Drivers/DriverException.h"
#include "Snapshot.h"
#include "schema.h"
#include <numeric>

namespace SSAGES
{

    class ParticleCoordinateCV : public CollectiveVariable
    {
    private:
        Label atomids_; 

        Dimension dim_;

    public:

        ParticleCoordinateCV(const Label& atomids, Dimension dim) : 
        atomids_(atomids), dim_(dim)
        {}


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

            auto n = atomids_.size();

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

            MPI_Allreduce(MPI_IN_PLACE, found.data(), n, MPI_INT, MPI_SUM, snapshot.GetCommunicator());
            unsigned ntot = std::accumulate(found.begin(), found.end(), 0, std::plus<int>());
            if(ntot != n)
                throw BuildException({
                    "ParticleCoordinateCV: Expected to find " + 
                    to_string(n) + 
                    " atoms, but only found " + 
                    to_string(ntot) + "."
                });         
        }


        void Evaluate(const Snapshot& snapshot) override
        {
            // Get local atom indices and compute COM. 
            std::vector<int> idx;
            snapshot.GetLocalIndices(atomids_, &idx);

            // Get data from snapshot. 
            auto n = snapshot.GetNumAtoms();
            const auto& masses = snapshot.GetMasses();

            // Initialize gradient.
            std::fill(grad_.begin(), grad_.end(), Vector3{0,0,0});
            grad_.resize(n, Vector3{0,0,0});
            
            // Compute total and center of mass.
            auto masstot = snapshot.TotalMass(idx);
            Vector3 com = snapshot.CenterOfMass(idx, masstot);

            // Assign CV value. 
            switch(dim_)
            {
                case Dimension::x:
                    val_ = com[0];
                    break;
                case Dimension::y:
                    val_ = com[1];
                    break;
                case Dimension::z:
                    val_ = com[2];
                    break;
            }

            // Assign gradient to appropriate atoms.
            for(auto& id : idx)
            {
                grad_[id][0] = (dim_ == Dimension::x) ? masses[id]/masstot : 0;
                grad_[id][1] = (dim_ == Dimension::y) ? masses[id]/masstot : 0;
                grad_[id][2] = (dim_ == Dimension::z) ? masses[id]/masstot : 0;
            }
        }


        double GetPeriodicValue(double Location) const override
        {
            return Location;
        }


        double GetDifference(const double Location) const override
        {
            return val_ - Location;
        }

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

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

            // Validate inputs.
            validator.Validate(json, path);
            if(validator.HasErrors())
                throw BuildException(validator.GetErrors());

            Label atomids;
            for(auto& id : json["atom_ids"])
                atomids.push_back(id.asInt());

            auto indextype = json.get("dimension","x").asString();

            Dimension index;
            if(indextype == "x")
                index = Dimension::x;
            else if(indextype == "y")
                index = Dimension::y;
            else if(indextype == "z")
                index = Dimension::z;
            else
                throw BuildException({"Could not obtain ParticleCoordinate dimension specified."});

            return new ParticleCoordinateCV(atomids, index);
        }
     };
}