WannierMinimizer.h

/*-------------------------------------------------------------------
Copyright 2014 Ravishankar Sundararaman

This file is part of JDFTx.

JDFTx is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

JDFTx is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with JDFTx.  If not, see <http://www.gnu.org/licenses/>.
-------------------------------------------------------------------*/

#ifndef JDFTX_WANNIER_WANNIERMINIMIZER_H
#define JDFTX_WANNIER_WANNIERMINIMIZER_H

#include <core/Util.h>
#include <core/Minimize.h>
#include <core/LatticeUtils.h>
#include <electronic/Everything.h>
#include <electronic/ColumnBundle.h>
#include <electronic/matrix.h>
#include <electronic/ColumnBundleTransform.h>
#include <wannier/Wannier.h>

struct WannierGradient : public std::vector<matrix>
{       const class WannierMinimizer* wmin;
    void init(const class WannierMinimizer* wmin);
        size_t ikStart() const;
        size_t ikStop() const;
};

//---- linear algebra functions required by Minimizable<WannierGradient> -----
WannierGradient clone(const WannierGradient& grad);
double dot(const WannierGradient& x, const WannierGradient& y);
WannierGradient& operator*=(WannierGradient& x, double alpha);
void axpy(double alpha, const WannierGradient& x, WannierGradient& y);
matrix randomMatrix(int nRows, int nCols);
void randomize(WannierGradient& x);

class WannierMinimizer : public Minimizable<WannierGradient>
{
public:
        WannierMinimizer(const Everything& e, const Wannier& wannier, bool needSuperOverride=false);
        virtual ~WannierMinimizer() {}
        void initTransformDependent(); //second half of common initialization that must happen after sub-class is fully initialized
        
        void saveMLWF(); //save wannier functions for all spins
        void saveMLWF(int iSpin); //save for specified spin
        
        //Interface for minimize:
        void step(const WannierGradient& grad, double alpha);
        double compute(WannierGradient* grad, WannierGradient* Kgrad); //identity preconditioner, but impose hermiticity constraint
        void constrain(WannierGradient& grad); //enforce hermiticity
        bool report(int iter);
        double sync(double x) const; 
        
        struct Kpoint : public QuantumNumber, public Supercell::KmeshTransform
        {       bool operator<(const Kpoint& other) const;
                bool operator==(const Kpoint& other) const;
        };
        
        struct KmeshEntry
        {       Kpoint point; 
                //State of system for Wannier minimize:
                int nIn; //number of bands that contribute to the Wannier subspace
                int nFixed; //number of bands that contribute fully to the Wannier subspace (cannot be partially mixed out)
                matrix U, Omega_U; 
                matrix U1; //Initial rotation into Wannier subspace (nBands x nIn)
                matrix U2; //Rotation within Wannier subspace (nCenters x nCenters)
        };

        //-------- Interface for subclasses that provide the objective function for Wannier minimization
        
        virtual void initialize(int iSpin)=0;
        
        virtual double getOmega(bool grad=false)=0;
        
        //Like getOmega, but for the subspace-invariant OmegaI instead.
        virtual double getOmegaI(bool grad=false)=0;
        
protected:
        friend struct WannierGradient;
        const Everything& e;
        const Wannier& wannier;
        const std::vector< matrix3<int> >& sym;
        int nCenters, nFrozen, nBands; 
        int nSpins, qCount; 
        int nSpinor; 
        std::vector<double> rSqExpect; 
        std::vector< vector3<> > rExpect; 
        std::vector<bool> pinned; 
        std::vector< vector3<> > rPinned; 
        
        //Supercell grid and basis:
        bool needSuper;
        GridInfo gInfoSuper;
        Basis basisSuper;
        QuantumNumber qnumSuper;
        std::map<vector3<int>,double> iCellMap, phononCellMap; //unit-cell indices in supercell (and weights to account for surface multiplicity)
        int nPhononModes; //number of phonon modes
        
        std::vector<KmeshEntry> kMesh; 
        std::set<Kpoint> kpoints; 
        std::shared_ptr<ColumnBundleTransform::BasisWrapper> basisWrapper, basisSuperWrapper; 
        std::map<Kpoint, std::shared_ptr<ColumnBundleTransform> > transformMap, transformMapSuper; 
        Basis basis; 
        
        //k-mesh MPI division:
        TaskDivision kDivision;
        size_t ikStart, ikStop;
        inline bool isMine(size_t ik) const { return kDivision.isMine(ik); }
        inline int whose(size_t ik) const { return kDivision.whose(ik); }

        //state MPI division (wrappers to ElecInfo)
        bool isMine_q(int ik, int iSpin) const { return e.eInfo.isMine(kMesh[ik].point.iReduced + iSpin*qCount); }
        int whose_q(int ik, int iSpin) const { return e.eInfo.whose(kMesh[ik].point.iReduced + iSpin*qCount); }
        
        ColumnBundle getWfns(const Kpoint& kpoint, int iSpin) const;
        std::vector<ColumnBundle> Cother; //wavefunctions from another process
        
        void axpyWfns(double alpha, const matrix& A, const Kpoint& kpoint, int iSpin, ColumnBundle& result) const;
        
        void axpyWfns_grad(double alpha, matrix& Omega_A, const Kpoint& kpoint, int iSpin, const ColumnBundle& Omega_result) const;
        
        ColumnBundle trialWfns(const Kpoint& kpoint) const;
        std::map< Kpoint, std::shared_ptr<ColumnBundle> > numericalOrbitals;
        
        matrix overlap(const ColumnBundle& C1, const ColumnBundle& C2) const;
        
        //Dump a named matrix variable to file, optionally zeroing out the real parts
        void dumpMatrix(const matrix& H, string varName, bool realPartOnly, int iSpin) const;
        
        static matrix fixUnitary(const matrix& U); //return an exactly unitary version of U (orthogonalize columns)
        
        //Preconditioner for Wannier optimization: identity by default, override in derived class to change
        virtual WannierGradient precondition(const WannierGradient& grad);
};

#endif // JDFTX_WANNIER_WANNIERMINIMIZER_H