ABF.cpp

 
#include "ABF.h"
#include "CVs/CVManager.h"
#include "Drivers/DriverException.h"
#include "Validator/ObjectRequirement.h"
#include "Snapshot.h"
#include "schema.h"
#include <algorithm>
#include <fstream>
#include <iostream>
 
using namespace Json;
// "Adaptive biasing force method for scalar and vector free energy calculations"
// Darve, Rodriguez-Gomez, Pohorille
// J. Chem. Phys. (2008)
namespace SSAGES
{
 
    // Pre-simulation hook.
    void ABF::PreSimulation(Snapshot* snapshot, const CVManager& cvmanager)
    {
        // Open/close outfile to create it fresh. 
        if(IsMasterRank(world_))
        {
            worldout_.open(filename_);
            worldout_.close();
        }
        
        // Convenience. Number of CVs.
        auto cvs = cvmanager.GetCVs(cvmask_);
        dim_ = cvs.size();
 
        // Size and initialize Fold_
        Fold_.setZero(dim_);
 
        // Initialize biases.
        biases_.resize(snapshot->GetPositions().size(), Vector3{0, 0, 0});
        
        // Initialize w \dot p's for finite difference. 
        wdotp1_.setZero(dim_);
        wdotp2_.setZero(dim_);  
    }
 
 
    void ABF::PostIntegration(Snapshot* snapshot, const CVManager& cvmanager)
    {
        
        ++iteration_;
 
        // Gather information.
        auto cvs = cvmanager.GetCVs(cvmask_);
        auto& vels = snapshot->GetVelocities();
        auto& mass = snapshot->GetMasses();
        auto& forces = snapshot->GetForces();
        auto n = snapshot->GetNumAtoms();
        
 
        //int coord = histCoords(cvs);
 
        Eigen::MatrixXd J(dim_, 3*n);
        
        std::vector<double> cvVals(dim_);
 
        // Fill J and CV. Each column represents grad(CV) with flattened Cartesian elements. 
        for(size_t i = 0; i < dim_; ++i)
        {
            auto& grad = cvs[i]->GetGradient();
            for(size_t j = 0; j < n; ++j)
                J.block<1, 3>(i,3*j) = grad[j];
            cvVals[i] = cvs[i]->GetValue();
        }
 
        //* Calculate W using Darve's approach (http://mc.stanford.edu/cgi-bin/images/0/06/Darve_2008.pdf).
        Eigen::MatrixXd Jmass = J.transpose();
        if(massweigh_)
        {
            for(size_t i = 0; i < forces.size(); ++i)
                Jmass.block(3*i, 0, 3, dim_) = Jmass.block(3*i, 0, 3, dim_)/mass[i];
        }
 
        Eigen::MatrixXd Minv = J*Jmass;
        MPI_Allreduce(MPI_IN_PLACE, Minv.data(), Minv.size(), MPI_DOUBLE, MPI_SUM, comm_);
        Eigen::MatrixXd Wt = Minv.inverse()*Jmass.transpose();  
        // Fill momenta.
        Eigen::VectorXd momenta(3*vels.size());
        for(size_t i = 0; i < vels.size(); ++i)
            momenta.segment<3>(3*i) = mass[i]*vels[i];
        
        // Compute dot(w,p)
        Eigen::VectorXd wdotp = Wt*momenta;
 
        
        // Reduce dot product across processors.
        MPI_Allreduce(MPI_IN_PLACE, wdotp.data(), wdotp.size(), MPI_DOUBLE, MPI_SUM, comm_);
        
    
        // Compute d(wdotp)/dt second order backwards finite difference. 
        // Adding old force removes bias. 
        Eigen::VectorXd dwdotpdt = unitconv_*(1.5*wdotp - 2.0*wdotp1_ + 0.5*wdotp2_)/timestep_ + Fold_;
 
        // If we are in bounds, sum force into running total.
        if(boundsCheck(cvVals) && IsMasterRank(comm_))
        {
            for(size_t i = 0; i < dim_; ++i)
                F_[i]->at(cvVals) += dwdotpdt[i];
            N_->at(cvVals)++;               
        }
        
        // Reduce data across processors.
        for(size_t i = 0; i < dim_; ++i)                
            MPI_Allreduce(F_[i]->data(), Fworld_[i]->data(), (F_[i]->size()), MPI_DOUBLE, MPI_SUM, world_); 
        MPI_Allreduce(N_->data(), Nworld_->data(), N_->size(), MPI_INT, MPI_SUM, world_);
 
        // If we are in bounds, store the old summed force.
        if(boundsCheck(cvVals))
        {
            for(size_t i=0; i < dim_; ++i)
                Fold_[i] = Fworld_[i]->at(cvVals)/std::max(min_, Nworld_->at(cvVals));
        }
    
        // Update finite difference time derivatives.
        wdotp2_ = wdotp1_;
        wdotp1_ = wdotp;        
 
        // Write out data to file.
        if(iteration_ % FBackupInterv_ == 0){
            WriteData();
        }
 
        // Calculate the bias from averaged F at current CV coordinates.
        // Or apply harmonic restraints to return CVs back in bounds.
        CalcBiasForce(snapshot, cvs, cvVals);       
 
        // Update the forces in snapshot by adding in the force bias from each
        // CV to each atom based on the gradient of the CV.
        for (size_t j = 0; j < forces.size(); ++j)
            forces[j] += biases_[j];    
    }
 
    // Post-simulation hook.
    void ABF::PostSimulation(Snapshot*, const CVManager&)
    {
        WriteData();
    }
 
    void ABF::CalcBiasForce(const Snapshot* snapshot, const CVList& cvs, const std::vector<double> &cvVals)
    {
        // Reset the bias.
        biases_.resize(snapshot->GetNumAtoms(), Vector3{0,0,0});
        for(auto& b : biases_)
            b.setZero();
        
        // Compute bias if within bounds.
        if(boundsCheck(cvVals))
        {
            for(size_t i = 0; i < dim_; ++i)
            {
                auto& grad = cvs[i]->GetGradient();
                for(size_t j = 0; j < biases_.size(); ++j)
                    biases_[j] -= (Fworld_[i]->at(cvVals))*grad[j]/std::max(min_,Nworld_->at(cvVals));
            }
        }
        // Otherwise, apply harmonic restraint if enabled.
        else
        {
            for(size_t i = 0; i < dim_; ++i)
            {
                double cvVal = cvs[i]->GetValue();
                auto k = 0.;
                auto x0 = 0.;
                
                if(isperiodic_[i])
                {
                    double periodsize = periodicboundaries_[i][1]-periodicboundaries_[i][0];
                    double cvRestrMidpoint = (restraint_[i][1]+restraint_[i][0])/2;
                    while((cvVal-cvRestrMidpoint) > periodsize/2)
                        cvVal -= periodsize;
                    while((cvVal-cvRestrMidpoint) < -periodsize/2)
                        cvVal += periodsize;
                }
 
                if(cvVal < restraint_[i][0])
                {
                    k = restraint_[i][2];
                    x0 = restraint_[i][0];
                }
                else if (cvVal > restraint_[i][1])
                {
                    k = restraint_[i][2];
                    x0 = restraint_[i][1];
                }
 
                auto& grad = cvs[i]->GetGradient();
                for(size_t j = 0; j < biases_.size(); ++j)
                    biases_[j] -= (cvVal - x0)*k*grad[j];
            }   
        }
    }
 
    // Write out the average generalized force.
    // Also write out Fworld and Nworld backups for restarts.
    void ABF::WriteData()
    {
        
        // Only one processor should be performing I/O.
        if(!IsMasterRank(world_))
            return;
 
        Nworld_->syncGrid();
 
        // Backup Fworld and Nworld.
        Nworld_->WriteToFile(Nworld_filename_);
        for(size_t i = 0 ; i < dim_; ++i)
        {
            Fworld_[i]->syncGrid();
            Fworld_[i]->WriteToFile(Fworld_filename_+std::to_string(i));
        }
 
        // Average the generalized force for each bin and print it out. 
        int gridPoints = 1;         
        for(size_t i = 0 ; i < dim_; ++i)
            gridPoints = gridPoints * N_->GetNumPoints(i);
 
        worldout_.open(filename_, std::ios::out);
        worldout_ << std::endl;
        worldout_ << "Iteration: " << iteration_ << std::endl;          
        worldout_ << "Printing out the current Adaptive Biasing Vector Field." << std::endl;
        worldout_ << "First (Nr of CVs) columns are the coordinates, the next (Nr of CVs) columns are components of the Adaptive Force vector at that point." << std::endl;
        worldout_ << "The columns are " << gridPoints << " long, mapping out a surface of ";
        
        for(size_t i = 0 ; i < dim_-1; ++i)
            worldout_ << (N_->GetNumPoints())[i] << " by " ;
        worldout_ << (N_->GetNumPoints(dim_-1)) << " points in " << dim_ << " dimensions." << std::endl;
        worldout_ << std::endl; 
        
        
        for(auto it = Nworld_->begin(); it != Nworld_->end(); ++it)
        {
            auto& val = *it; 
            auto coord = it.coordinates(); 
            for(auto& c : coord)
                worldout_ << std::setw(14) << std::setprecision(8) << std::fixed << c << " ";
            for(size_t i = 0 ; i < dim_; ++i)
                worldout_ << std::setw(14) << std::setprecision(8) << std::fixed << Fworld_[i]->at(coord)/std::max(val,min_) << " ";
            worldout_ << std::endl;
        }
    
        worldout_ << std::endl;
        worldout_ << std::endl;
        worldout_.close();
    }
    
    // Checks whether walker is within CV bounds.
    bool ABF::boundsCheck(const std::vector<double> &CVs)
    {
        for(size_t i = 0; i < dim_ ; ++i)
        {
            if((CVs[i] < N_->GetLower(i)) || (CVs[i] > N_->GetUpper(i)))
            {
                return false;
            }
        }
        return true;
    }
 
 
    ABF* ABF::Build(const Value& json, 
                        const MPI_Comm& world,
                        const MPI_Comm& comm,
                        const std::string& path)
    {
        ObjectRequirement validator;
        Value schema;
        CharReaderBuilder rbuilder;
        CharReader* reader = rbuilder.newCharReader();
 
        reader->parse(JsonSchema::ABFMethod.c_str(),
                      JsonSchema::ABFMethod.c_str() + JsonSchema::ABFMethod.size(),
                      &schema, nullptr);
        validator.Parse(schema, path);
 
        // Validate inputs.
        validator.Validate(json, path);
        if(validator.HasErrors())
            throw BuildException(validator.GetErrors());
 
        //bool multiwalkerinput = false;
 
        // Obtain lower bound for each CV.
        std::vector<double> minsCV;
        for(auto& mins : json["CV_lower_bounds"])
            if(mins.isArray())
                {
                throw BuildException({"Separate inputs for multi-walker not fully implemented. Please use global entries for CV ranges"});
                //multiwalkerinput = true;
                //for(auto& bound : mins)
                //  minsCV.push_back(bound.asDouble());
                }
            else
                minsCV.push_back(mins.asDouble());
            
        // Obtain upper bound for each CV.
        std::vector<double> maxsCV;
        for(auto& maxs : json["CV_upper_bounds"])
            /*if(maxs.isArray())
                {
                for(auto& bound : maxs)
                    maxsCV.push_back(bound.asDouble());
                }*/
            //else
                maxsCV.push_back(maxs.asDouble());
 
        // Obtain number of bins for each CV dimension.
        std::vector<int> binsCV;
        for(auto& bins : json["CV_bins"])
            binsCV.push_back(bins.asInt());
 
        // Obtain lower bounds for restraints for each CV.
        std::vector<double> minsrestCV;
        for(auto& mins : json["CV_restraint_minimums"])
            /*if(mins.isArray())
                {
                for(auto& bound : mins)
                    minsrestCV.push_back(bound.asDouble());
                }*/
            //else
                minsrestCV.push_back(mins.asDouble());
 
        // Obtain upper bounds for restraints for each CV.      
        std::vector<double> maxsrestCV;
        for(auto& maxs : json["CV_restraint_maximums"])
            /*if(maxs.isArray())
                {
                for(auto& bound : maxs)
                    maxsrestCV.push_back(bound.asDouble());
                }*/
            //else
                maxsrestCV.push_back(maxs.asDouble());
        
        // Obtain harmonic spring constant for restraints for each CV.
        std::vector<double> springkrestCV;
        for(auto& springkrest : json["CV_restraint_spring_constants"])
        {
            if(springkrest.asDouble() < 0)
            {
                throw BuildException({"Restraint spring constants must be non-negative."});
            }
            springkrestCV.push_back(springkrest.asDouble());
        }
 
        // Obtain periodicity information for restraints for each CV for the purpose of correctly applying restraints through periodic boundaries.
        std::vector<bool> isperiodic;
        for(auto& isperCV : json["CV_isperiodic"])
            isperiodic.push_back(isperCV.asBool());
        
        // Obtain lower periodic boundary for each CV.
        std::vector<double> minperboundaryCV;
        for(auto& minsperCV : json["CV_periodic_boundary_lower_bounds"])
            minperboundaryCV.push_back(minsperCV.asDouble());
 
        // Obtain upper periodic boundary for each CV.
        std::vector<double> maxperboundaryCV;
        for(auto& maxsperCV : json["CV_periodic_boundary_upper_bounds"])
            maxperboundaryCV.push_back(maxsperCV.asDouble());
 
        // Verify inputs are all correct lengths.
        if(!((   minsCV.size() == maxsCV.size() && 
             maxsCV.size() == minsrestCV.size() &&
             minsrestCV.size() == maxsrestCV.size()) &&
            (binsCV.size() == springkrestCV.size() &&
             springkrestCV.size() == isperiodic.size())))       
            throw BuildException({"CV lower bounds, upper bounds, restrain minimums, restrains maximums must match in size. Bins, spring constants and periodicity info must match in size."});
 
        // Verify that all periodicity information is provided if CVs are periodic.
        bool anyperiodic=false;
        for(size_t i = 0; i<isperiodic.size(); ++i)
            if(isperiodic[i])
            {
                anyperiodic = true;
                if(!(isperiodic.size() == minperboundaryCV.size() &&
                     minperboundaryCV.size() == maxperboundaryCV.size()))
                    throw BuildException({"If any CV is defined as periodic, please define the full upper and lower bound vectors. They should both have the same number of entries as CV lower bounds, upper bounds... Entries corresponding to non-periodic CVs will not be used."});
            }
 
        size_t dim = binsCV.size();
        
        // Read in JSON info.
        auto FBackupInterv = json.get("output_frequency", 1000).asInt();
        auto unitconv = json.get("unit_conversion", 0).asDouble();
        auto timestep = json.get("timestep", 2).asDouble();
        auto min = json.get("minimum_count", 200).asDouble();
        auto massweigh = json.get("mass_weighing",false).asBool();          
 
        std::vector<std::vector<double>> histdetails;
        std::vector<std::vector<double>> restraint;
        std::vector<std::vector<double>> periodicboundaries;
        std::vector<double> temp1(3);
        std::vector<double> temp2(3);
        std::vector<double> temp3(2);
 
        auto freq = json.get("frequency", 1).asInt();
        auto filename = json.get("output_file", "F_out").asString();
        auto Nworld_filename = json.get("Nworld_output_file", "Nworld").asString();
        auto Fworld_filename = json.get("Fworld_output_file", "Fworld_cv").asString();
 
        // Generate the grids based on JSON.
        Grid<int> *N;
        Grid<int> *Nworld;
        std::vector<Grid<double>*> F(dim);
        std::vector<Grid<double>*> Fworld(dim);
 
        //unsigned int wid = GetWalkerID(world, comm);
 
        // This feature is disabled for now.
 
        /*if(multiwalkerinput)
        {
            for(size_t i=0; i<dim; ++i)
            {
                temp1 = {minsCV[i+wid*dim], maxsCV[i+wid*dim], binsCV[i]};
                temp2 = {minsrestCV[i+wid*dim], maxsrestCV[i+wid*dim], springkrestCV[i]};
                histdetails.push_back(temp1);
                restraint.push_back(temp2);
                if(anyperiodic)
                {
                    temp3 = {minperboundaryCV[i], maxperboundaryCV[i]};
                    periodicboundaries.push_back(temp3);
                }
            }
            for(size_t i=0; i<dim; ++i)
            {
                minsCVperwalker[i] = histdetails[i][0];
                maxsCVperwalker[i] = histdetails[i][1];
            }
 
            N= new Grid<int>(binsCV, minsCVperwalker, maxsCVperwalker, isperiodic);
            Nworld= new Grid<int>(binsCV, minsCVperwalker, maxsCVperwalker, isperiodic);
            
            for(auto& grid : F)
            {
                grid= new Grid<double>(binsCV, minsCVperwalker, maxsCVperwalker, isperiodic);
            }
            for(auto& grid : Fworld)
            {
                grid= new Grid<double>(binsCV, minsCVperwalker, maxsCVperwalker, isperiodic);
            }
            
        }
        */
        //else
        //{
        
        // Appropriately size the grids.
        
            N = new Grid<int>(binsCV, minsCV, maxsCV, isperiodic);
            Nworld = new Grid<int>(binsCV, minsCV, maxsCV, isperiodic);
            
            for(auto& grid : F)
            {
                grid = new Grid<double>(binsCV, minsCV, maxsCV, isperiodic);
            }
            for(auto& grid : Fworld)
            {
                grid = new Grid<double>(binsCV, minsCV, maxsCV, isperiodic);
            }
 
            for(size_t i = 0; i < dim; ++i)
            {
                temp1 = {minsCV[i], maxsCV[i], double(binsCV[i])};
                temp2 = {minsrestCV[i], maxsrestCV[i], springkrestCV[i]};
                histdetails.push_back(temp1);
                restraint.push_back(temp2);
                if(anyperiodic)
                {
                    temp3 = {minperboundaryCV[i], maxperboundaryCV[i]};
                    periodicboundaries.push_back(temp3);
                }
            }
        //}
 
        // Check if a restart is requested.
        bool restart = json.get("restart",false).asBool();
 
        // Either load previous grid, or back it up.
        if(IsMasterRank(world))
        {
            if(restart)
            {
                std::cout << "Attempting to load data from a previous run of ABF..." << std::endl;
                N->LoadFromFile(Nworld_filename);
                for(size_t i = 0; i < dim; ++i)
                {
                    F[i]->LoadFromFile(Fworld_filename+std::to_string(i));
                }
            }
            else
            {
                std::ifstream Nworldbackupsource(Nworld_filename, std::ios::binary);
                if(Nworldbackupsource)
                {
                    std::cout << "Backing up previous copy of Nworld." << std::endl;
                    std::ofstream Nworldbackuptarget(Nworld_filename+"_backup", std::ios::binary);
                    Nworldbackuptarget << Nworldbackupsource.rdbuf();
                }
                for(size_t i = 0; i < dim; ++i)
                {
                    std::ifstream Fworldbackupsource(Fworld_filename+std::to_string(i), std::ios::binary);
                    if(Fworldbackupsource)
                    {
                        std::cout << "Backing up previous copy of Fworld_cv"+std::to_string(i)+"." << std::endl;
                        std::ofstream Fworldbackuptarget(Fworld_filename+std::to_string(i)+"_backup", std::ios::binary);
                        Fworldbackuptarget << Fworldbackupsource.rdbuf();
                    }
                }
            }
        }
 
        auto* m = new ABF(world, comm, N, Nworld, F, Fworld, restraint, isperiodic, periodicboundaries, min, massweigh, filename, Nworld_filename, Fworld_filename, histdetails, FBackupInterv, unitconv, timestep, freq);
 
        if(json.isMember("iteration"))
            m->SetIteration(json.get("iteration",0).asInt());
 
        return m;
 
    }
 
}