/home/js850/projects/pele/source/pele/frozen_atoms.h

#ifndef _PELE_FROZEN_ATOMS_H
#define _PELE_FROZEN_ATOMS_H

#include <vector>
#include <algorithm>
#include <set>
#include <assert.h>
#include <iostream>
#include <memory>

#include "array.h"
#include "base_potential.h"

namespace pele{
class FrozenCoordsConverter
{
private:
    std::vector<double> const _reference_coords;
    std::vector<size_t> _frozen_dof;
    std::vector<size_t> _mobile_dof;
public:
    FrozenCoordsConverter(Array<double> const & reference_coords, 
            Array<size_t> const & frozen_dof) :
        _reference_coords(reference_coords.begin(), reference_coords.end())
    {
        //populate _frozen_dof after removing duplicates and sorting
        std::set<size_t> frozenset(frozen_dof.begin(), frozen_dof.end());
        _frozen_dof = std::vector<size_t>(frozenset.begin(), frozenset.end());
        std::sort(_frozen_dof.begin(), _frozen_dof.end());

        // do a sanity check
        if (_frozen_dof.size() > 0){
            if (_frozen_dof[_frozen_dof.size()-1] >= (size_t)ndof()) {
                throw std::runtime_error("index of frozen degree of freedom is out of bounds");
            }
        }

        //populate _mobile_dof
        _mobile_dof = std::vector<size_t>(ndof() - ndof_frozen());
        size_t imobile = 0;
        for (size_t i=0; i<_reference_coords.size(); ++i){
            // if degree of freedom i is not in frozen, add it to _mobile_dof
            if (frozenset.count(i) == 0){
                _mobile_dof[imobile] = i;
                ++imobile;
            }
        }
        assert(imobile == _mobile_dof.size());

    }

    size_t ndof() const { return _reference_coords.size(); }
    size_t ndof_frozen() const { return _frozen_dof.size(); }
    size_t ndof_mobile() const { return _mobile_dof.size(); }

    Array<double> get_reduced_coords(Array<double> const &full_coords){
        if (full_coords.size() != ndof()) {
            std::invalid_argument("full_coords has the wrong size");
        }
        Array<double> reduced_coords(ndof_mobile());
        for (size_t i=0; i < ndof_mobile(); ++i){
            reduced_coords[i] = full_coords[_mobile_dof[i]];
        }
        return reduced_coords;
    }

    Array<double> get_full_coords(Array<double> const &reduced_coords)
    {
        if (reduced_coords.size() != ndof_mobile()) {
            std::invalid_argument("reduced_coords has the wrong size");
        }
        // make a new array full_coords as a copy of _reference_coords.
        //Array<double> const a(_reference_coords); //wrap _reference_coords in an Array (returns error due to _reference_coords being const)
        Array<double> full_coords(ndof());
        for (size_t i=0; i < ndof(); ++i){
            full_coords[i] = _reference_coords[i];
        }
        // replace the mobile degrees of freedom with those in reduced_coords
        for (size_t i=0; i < _mobile_dof.size(); ++i){
            full_coords[_mobile_dof[i]] = reduced_coords[i];
        }
        return full_coords;
    }

    Array<double> get_full_grad(Array<double> const &reduced_grad)
    {
        if (reduced_grad.size() != ndof_mobile()) {
            std::invalid_argument("reduced_grad has the wrong size");
        }
        Array<double> full_grad(ndof(), 0.);
        // replace the mobile degrees of freedom with those in reduced_grad
        for (size_t i=0; i < _mobile_dof.size(); ++i){
            full_grad[_mobile_dof[i]] = reduced_grad[i];
        }
        return full_grad;
    }

    Array<double> get_reduced_hessian(Array<double> const full_hess)
    {
        if (full_hess.size() != ndof()*ndof()) {
            std::invalid_argument("full_hessian has the wrong size");
        }
        Array<double> reduced_hess(ndof_mobile() * ndof_mobile());
//                size_t const N = ndof_mobile();
        for (size_t i=0; i < ndof_mobile(); ++i){
            for (size_t j=0; j < ndof_mobile(); ++j){
                size_t k_red = i * ndof_mobile() + j;
                size_t k_full = _mobile_dof[i] * ndof() + _mobile_dof[j];
                reduced_hess[k_red] = full_hess[k_full];
            }
        }
        return reduced_hess;
    }

};


template<typename PotentialType>
class FrozenPotentialWrapper : public BasePotential
{
public:
    FrozenCoordsConverter coords_converter;
protected:
    std::shared_ptr<PotentialType> _underlying_potential;

    FrozenPotentialWrapper(std::shared_ptr<PotentialType> potential,
            Array<double> const &reference_coords,
            Array<size_t> const & frozen_dof) :
        coords_converter(reference_coords, frozen_dof),
        _underlying_potential(potential)
    {}

public:
    ~FrozenPotentialWrapper() {}

//            inline size_t ndof() const { return coords_converter.ndof(); }
//            inline size_t ndof_frozen() const { return coords_converter.ndof_frozen(); }
//            inline size_t ndof_mobile() const { return coords_converter.ndof_mobile(); }
//            inline Array<double> get_reduced_coords(Array<double> const &full_coords){
//                return coords_converter.get_reduced_coords(full_coords);
//            }
//            Array<double> get_full_coords(Array<double> const &reduced_coords){
//                return coords_converter.get_full_coords(reduced_coords);
//            }

    inline double get_energy(Array<double> reduced_coords) 
    {
        if (reduced_coords.size() != coords_converter.ndof_mobile()){
            throw std::runtime_error("reduced coords does not have the right size");
        }
        Array<double> full_coords(coords_converter.get_full_coords(reduced_coords));
        return _underlying_potential->get_energy(full_coords);
    }

    inline double get_energy_gradient(Array<double> reduced_coords, Array<double> reduced_grad) 
    {
        if (reduced_coords.size() != coords_converter.ndof_mobile()){
            throw std::runtime_error("reduced coords does not have the right size");
        }
        if (reduced_grad.size() != coords_converter.ndof_mobile()) {
            throw std::invalid_argument("reduced_grad has the wrong size");
        }
        
        Array<double> full_coords(coords_converter.get_full_coords(reduced_coords));
        Array<double> gfull(coords_converter.ndof());
        double energy = _underlying_potential->get_energy_gradient(full_coords, gfull);
        Array<double> gred = coords_converter.get_reduced_coords(gfull);
        for (size_t i = 0; i < gred.size(); ++i){
            reduced_grad[i] = gred[i];
        }
        return energy;
    }

    inline double get_energy_gradient_hessian(Array<double> reduced_coords,
            Array<double> reduced_grad, Array<double> reduced_hess)
    {
        if (reduced_coords.size() != coords_converter.ndof_mobile()){
            throw std::runtime_error("reduced coords does not have the right size");
        }
        if (reduced_grad.size() != coords_converter.ndof_mobile()) {
            throw std::invalid_argument("reduced_grad has the wrong size");
        }
        if (reduced_hess.size() != coords_converter.ndof_mobile()*coords_converter.ndof_mobile()){
            throw std::invalid_argument("reduced_hess has the wrong size");
        }
        Array<double> full_coords(coords_converter.get_full_coords(reduced_coords));
        Array<double> gfull(coords_converter.ndof());
        Array<double> hfull(coords_converter.ndof()*coords_converter.ndof());
        const double energy = _underlying_potential->get_energy_gradient_hessian(full_coords, gfull, hfull);
        Array<double> gred = coords_converter.get_reduced_coords(gfull);
        Array<double> hred = coords_converter.get_reduced_hessian(hfull);
        reduced_grad.assign(gred);
        reduced_hess.assign(hred);
        return energy;
    }
};
}

#endif