AngleCV.h

#pragma once 

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

#include <array>
#include <cmath>

namespace SSAGES
{
    class AngleCV : public CollectiveVariable
    {
    private:
        Label atomids_;

    public:


        AngleCV(int atomid1, int atomid2, int atomid3) :
        atomids_({atomid1, atomid2, atomid3})
        {
            bounds_ = {{0,M_PI}};
        }


        void Initialize(const Snapshot& snapshot) override
        {
            using std::to_string;
            
            std::vector<int> found;
            snapshot.GetLocalIndices(atomids_, &found);
            int nfound = found.size();
            MPI_Allreduce(MPI_IN_PLACE, &nfound, 1, MPI_INT, MPI_SUM, snapshot.GetCommunicator());
            
            if(nfound != 3)
                throw BuildException({
                    "TorsionalCV: Expected to find " + 
                    to_string(3) + 
                    " atoms, but only found " + 
                    to_string(nfound) + "."
                }); 
        }


        void Evaluate(const Snapshot& snapshot) override
        {
            // Get data from snapshot. 
            auto n = snapshot.GetNumAtoms();
            const auto& pos = 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});

            Vector3 xi{0, 0, 0}, xj{0, 0, 0}, xk{0, 0, 0};

            auto iindex = snapshot.GetLocalIndex(atomids_[0]);
            auto jindex = snapshot.GetLocalIndex(atomids_[1]);
            auto kindex = snapshot.GetLocalIndex(atomids_[2]);

            if(iindex != -1) xi = pos[iindex];
            if(jindex != -1) xj = pos[jindex];
            if(kindex != -1) xk = pos[kindex];

            // By performing a reduce, we actually collect all. This can 
            // be converted to a more intelligent allgater on rank then bcast. 
            MPI_Allreduce(MPI_IN_PLACE, xi.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            MPI_Allreduce(MPI_IN_PLACE, xj.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            MPI_Allreduce(MPI_IN_PLACE, xk.data(), 3, MPI_DOUBLE, MPI_SUM, comm);

            // Two vectors
            auto rij = snapshot.ApplyMinimumImage(xi - xj);
            auto rkj = snapshot.ApplyMinimumImage(xk - xj);

            auto dotP = rij.dot(rkj);
            auto nrij = rij.norm();
            auto nrkj = rkj.norm();

            val_ = acos(dotP/(nrij*nrkj));

            // Calculate gradients
            auto prefactor = -1.0/(sqrt(1. - dotP/(nrij*nrkj))*nrij*nrkj);

            Vector3 gradi{0, 0, 0}, gradk{0, 0, 0};
            if(iindex != -1) gradi = prefactor*(rkj - dotP*rij/(nrij*nrij));
            if(kindex != -1) gradk = prefactor*(rij - dotP*rkj/(nrkj*nrkj));
            MPI_Allreduce(MPI_IN_PLACE, gradi.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            MPI_Allreduce(MPI_IN_PLACE, gradk.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            if(iindex != -1) grad_[iindex] = gradi; 
            if(kindex != -1) grad_[kindex] = gradk;
            if(jindex != -1) grad_[jindex] = -gradi - gradk;
        }

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

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

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

            return new AngleCV(atomids[0], atomids[1], atomids[2]);
        }
    };
}