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

        bool periodic_;

    public:

        TorsionalCV(int atomid1, int atomid2, int atomid3, int atomid4, bool periodic) : 
        atomids_({atomid1, atomid2, atomid3, atomid4}), periodic_(periodic)
        {
        }


        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 != 4)
                throw BuildException({
                    "TorsionalCV: Expected to find " + 
                    to_string(4) + 
                    " 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}, xl{0, 0, 0};

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

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

            // 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);
            MPI_Allreduce(MPI_IN_PLACE, xl.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            
            //Calculate pertinent vectors
            auto F = snapshot.ApplyMinimumImage(xi - xj);
            auto G = snapshot.ApplyMinimumImage(xj - xk);
            auto H = snapshot.ApplyMinimumImage(xl - xk);
            auto A = F.cross(G);
            auto B = H.cross(G);
            
            auto y = B.cross(A).dot(G.normalized());
            auto x = A.dot(B);

            val_ = atan2(y, x);

            auto Zed = F.dot(G.normalized())/A.dot(A); 
            auto Ned = H.dot(G.normalized())/B.dot(B);

            Vector3 gradi{0,0,0}, gradl{0,0,0}; 
            if(iindex != -1) gradi = -G.norm()*A/A.dot(A);
            if(lindex != -1) gradl = G.norm()*B/B.dot(B);
            MPI_Allreduce(MPI_IN_PLACE, gradi.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            MPI_Allreduce(MPI_IN_PLACE, gradl.data(), 3, MPI_DOUBLE, MPI_SUM, comm);
            if(iindex != -1) grad_[iindex] = gradi;
            if(lindex != -1) grad_[lindex] = gradl;
            if(jindex != -1) grad_[jindex] = Zed*A - Ned*B - gradi;
            if(kindex != -1) grad_[kindex] = Ned*B  - Zed*A - gradl;
        }


        double GetPeriodicValue(double Location) const override
        {
            if(!periodic_)
                return Location;

            int n = (int)(Location/(2.0*M_PI));
            double PeriodicLocation;

            PeriodicLocation = Location - n*M_PI;
            if(PeriodicLocation < -M_PI)
                PeriodicLocation += 2.0*M_PI;
            else if (PeriodicLocation > M_PI)
                PeriodicLocation -= 2.0*M_PI;

            return PeriodicLocation;
        }


        double GetDifference(const double Location) const override
        {
            double PeriodicDiff = val_ - Location;

            if(!periodic_)
                return PeriodicDiff;

            PeriodicDiff = GetPeriodicValue(PeriodicDiff);

            if(PeriodicDiff > M_PI)
                PeriodicDiff -= 2.0*M_PI;
            else if(PeriodicDiff < -M_PI)
                PeriodicDiff += 2.0*M_PI;

            return PeriodicDiff;
        }


        double GetMinimumImage(const double Location) const override
        {
            double PeriodicDiff = val_ - Location;

            if(!periodic_)
                return val_;    

            if(PeriodicDiff > M_PI)
                return (val_ - 2.0*M_PI);
            else if(PeriodicDiff < -M_PI)
                return (val_ + 2.0*M_PI);   

            return val_;
        }

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

            reader.parse(JsonSchema::TorsionalCV, 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());

            auto periodic = json.get("periodic", true).asBool();
            return new TorsionalCV(atomids[0], atomids[1], atomids[2], atomids[3], periodic);
        }
    };
}