CFF.h

 
#pragma once
 
#include "Method.h"
#include "Grids/Grid.h"
#include "nnet/nnet.h"
 
namespace SSAGES
{
 
    class CFF : public Method
    {
    private:
        Eigen::VectorXi topol_;
 
        unsigned int sweep_, nsweep_;
 
        unsigned int citers_;
 
        nnet::neural_net net_;
 
        nnet::neural_net net2_;
 
        double timestep_;
 
        double pweight_, weight_;
 
        double temp_, kbt_;
 
        std::vector<Grid<double>*> F_;
 
        std::vector<Grid<double>*> Fworld_;
 
        Eigen::VectorXd Fold_;
 
        Eigen::VectorXd wdotp1_;
 
        Eigen::VectorXd wdotp2_;
 
        Grid<Eigen::VectorXd>* fgrid_;
 
        Grid<unsigned int>* hgrid_;
 
        Eigen::ArrayXi Nworld_;
 
        Grid<double>* ugrid_;
 
        double ratio_;
 
        Eigen::MatrixXd hist_, bias_;
 
        std::vector<double> lowerb_, upperb_;
 
        std::vector<double> lowerk_, upperk_;
 
        std::string outfile_;
 
        bool overwrite_;
 
        void TrainNetwork();
 
        void WriteBias();
 
        int dim_;
 
        double unitconv_;
 
        int min_;
 
        Eigen::MatrixXd force_to_val_ratio_;
 
    public:
 
        CFF(
            const MPI_Comm& world,
            const MPI_Comm& comm,
            const Eigen::VectorXi& topol,
            Grid<Eigen::VectorXd>* fgrid,
            Grid<unsigned int>* hgrid,
            Grid<double>* ugrid,
            std::vector<Grid<double>*> F,
            std::vector<Grid<double>*> Fworld,
            const std::vector<double>& lowerb,
            const std::vector<double>& upperb,
            const std::vector<double>& lowerk,
            const std::vector<double>& upperk,
            double temperature,
            double unitconv,
            double timestep,
            double weight,
            unsigned int nsweep,
            int min
        ) :
            Method(1, world, comm), topol_(topol), sweep_(0), nsweep_(nsweep), citers_(0),
            net_(topol), net2_(topol), timestep_(timestep), pweight_(1.), weight_(weight),
            temp_(temperature), kbt_(), F_(F), Fworld_(Fworld), fgrid_(fgrid),
            hgrid_(hgrid), ugrid_(ugrid), hist_(), bias_(), lowerb_(lowerb),
            upperb_(upperb), lowerk_(lowerk), upperk_(upperk), outfile_("CFF.out"),
            overwrite_(true), unitconv_(unitconv), min_(min)
        {
            // Create histogram grid matrix.
            hist_.resize(hgrid_->size(), hgrid_->GetDimension());
 
            // Fill it up.
            int i = 0;
            for(auto it = hgrid_->begin(); it != hgrid_->end(); ++it)
            {
                auto coord = it.coordinates();
                auto n = coord.size();
                for(decltype(n) j = 0; j < n; ++j)
                    hist_(i, j) = coord[j];
                ++i;
            }
 
            // Initialize free energy surface vector.
            bias_.resize(hgrid_->size(), 1);
            net_.forward_pass(hist_);
            net2_.forward_pass(hist_);
            bias_.array() = net_.get_activation().col(0).array();
        }
 
        void PreSimulation(Snapshot* snapshot, const class CVManager& cvmanager) override;
 
        void PostIntegration(Snapshot* snapshot, const class CVManager& cvmanager) override;
 
        void PostSimulation(Snapshot* snapshot, const class CVManager& cvmanager) override;
 
        void SetPrevWeight(double h)
        {
            pweight_ = h;
        }
 
        void SetOutput(const std::string& outfile)
        {
            outfile_ = outfile;
        }
 
        void SetOutputOverwrite(bool overwrite)
        {
            overwrite_ = overwrite;
        }
 
        void SetConvergeIters(unsigned int citers)
        {
            citers_ = citers;
        }
 
        void SetMaxIters(unsigned int iters)
        {
            auto params = net_.get_train_params();
            params.max_iter = iters;
            net_.set_train_params(params);
 
            auto params2 = net2_.get_train_params();
            params2.max_iter = iters;
            net2_.set_train_params(params2);
        }
 
        void SetMinLoss(double loss)
        {
            auto params = net_.get_train_params();
            params.min_loss = loss;
            net_.set_train_params(params);
 
            auto params2 = net2_.get_train_params();
            params2.min_loss = loss;
            net2_.set_train_params(params2);
        }
 
        void SetRatio(double trainratio)
        {
            auto params = net_.get_train_params();
            params.ratio = trainratio;
            net_.set_train_params(params);
 
            auto params2 = net2_.get_train_params();
            params2.ratio = trainratio;
            net2_.set_train_params(params2);
        }
 
        static CFF* Build(
            const Json::Value& json,
            const MPI_Comm& world,
            const MPI_Comm& comm,
            const std::string& path
        );
 
        ~CFF()
        {
            delete fgrid_;
            delete hgrid_;
        }
    };
}