ParticlePositionCV.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 ParticlePositionCV : public CollectiveVariable
    {
    private:
        Label atomids_; 

        Vector3 point_;

        Bool3 fix_; 
    
    public:

        ParticlePositionCV(const Label& atomids, const Vector3& position, bool fixx, bool fixy, bool fixz) : 
        atomids_(atomids), point_(position), fix_{fixx, fixy, fixz}
        {
        }


        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({
                    "ParticlePositionCV: 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);

            // Compute difference between point and account for requested 
            // dimensions only.
            Vector3 dx = snapshot.ApplyMinimumImage(com - point_).cwiseProduct(fix_.cast<double>());
            val_ = dx.norm();

            // If distance is zero, we have nothing to do.
            if(val_ == 0)
                return;

            for(auto& id : idx)
                grad_[id] = dx/val_*masses[id]/masstot;
        }
        
        static ParticlePositionCV* Build(const Json::Value& json, const std::string& path)
        {
            Json::ObjectRequirement validator;
            Json::Value schema;
            Json::Reader reader;

            reader.parse(JsonSchema::ParticlePositionCV, 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());
            
            Vector3 position;
            position[0] = json["position"][0].asDouble();
            position[1] = json["position"][1].asDouble();
            position[2] = json["position"][2].asDouble();

            auto fixx = json["fix"][0].asBool();
            auto fixy = json["fix"][1].asBool();
            auto fixz = json["fix"][2].asBool();

            return new ParticlePositionCV(atomids, position, fixx, fixy, fixz);
        }
    };
}