matrix.h

/*-------------------------------------------------------------------
Copyright 2011 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_ELECTRONIC_MATRIX_H
#define JDFTX_ELECTRONIC_MATRIX_H

#include <electronic/common.h>
#include <core/ManagedMemory.h>
#include <core/matrix3.h>
#include <core/scaled.h>
#include <core/GridInfo.h>
#include <gsl/gsl_cblas.h>

class diagMatrix : public std::vector<double>
{       
public:
        diagMatrix(int N=0, double d=0.) : std::vector<double>(N,d) {}
        int nRows() const { return size(); }
        int nCols() const { return size(); }
        bool isScalar(double absTol=1e-14, double relTol=1e-14) const; 
        
        diagMatrix& operator*=(double s) { for(double& d: *this) d*=s; return *this; }
        
        //Splicing operations:
        diagMatrix operator()(int iStart, int iStop) const; 
        diagMatrix operator()(int iStart, int iStep, int iStop) const; 
        void set(int iStart, int iStop, const diagMatrix& m); 
        void set(int iStart, int iStep, int iStop, const diagMatrix& m); 

        void scan(FILE* fp); 
        void print(FILE* fp, const char* fmt="%lg\t") const; 

        //Inter-process communication:
        void send(int dest, int tag=0) const; //send to another process
        void recv(int src, int tag=0); //receive from another process
        void bcast(int root=0); //synchronize across processes (using value on specified root process)
        void allReduce(MPIUtil::ReduceOp op, bool safeMode=false); //apply all-to-all reduction (see MPIUtil::allReduce)
};

class matrix : public ManagedMemory
{
        int nr; 
        int nc; 

public:
        int nRows() const { return nr; }
        int nCols() const { return nc; }
        int index(int i, int j) const { return nr*j+i; } 
        explicit operator bool() const { return nr*nc; } 

        void init(int nrows,int ncols, bool onGpu=false); 
        void reshape(int nrows, int ncols); 
        matrix(int nrows=0, int ncols=0, bool onGpu=false);
        matrix(const matrix& m1); 
        matrix(matrix&& m1); 
        matrix(const diagMatrix&); 
        matrix(const std::vector<complex>&); 
        explicit matrix(const matrix3<>&); 
        
        matrix& operator=(const matrix& m1); 
        matrix& operator=(matrix&& m1); 

        complex operator()(int i, int j) const;
        
        matrix operator()(int iStart, int iStop, int jStart, int jStop) const { return (*this)(iStart,1,iStop, jStart,1,jStop); }
        
        matrix operator()(int iStart, int iStep, int iStop, int jStart, int jStep, int jStop) const;

        void set(int i, int j, complex m);
        
        void set(int iStart, int iStop, int jStart, int jStop, const matrix& m) { set(iStart,1,iStop, jStart,1,jStop, m); }
        
        void set(int iStart, int iStep, int iStop, int jStart, int jStep, int jStop, const matrix& m);
        
        void scan(FILE* fp, const char* fmt="%lg%+lgi"); 
        void scan_real(FILE* fp); 
        void print(FILE* fp, const char* fmt="%lg%+lgi\t") const; 
        void print_real(FILE* fp, const char* fmt="%lg\t") const; 
        
        void diagonalize(matrix& evecs, diagMatrix& eigs) const; 
        void diagonalize(matrix& levecs, std::vector<complex>& eigs, matrix& revecs) const; 
        void svd(matrix& U, diagMatrix& S, matrix& Vdag) const; 
        
        //---------octave-like slicing operators on scalar fields converted to Nx*Ny*Nz x 1 matrices --------------
        complex getElement(vector3< int > index, GridInfo& gInfo);
        matrix getLine(vector3<int> line, vector3<int> point, GridInfo& gInfo);
        matrix getPlane(vector3<int> normal, vector3<int> point, GridInfo& gInfo);
};

struct matrixScaledTransOp
{       const matrix& mat; 
        double scale; 
        CBLAS_TRANSPOSE op; 
        
        int nRows() const { return op==CblasNoTrans ? mat.nRows() : mat.nCols(); }
        int nCols() const { return op==CblasNoTrans ? mat.nCols() : mat.nRows(); }
        complex conjOp(complex a) const { return (op==CblasConjTrans) ? a.conj() : a; } 
        int index(int i, int j) const { return op==CblasNoTrans ? mat.index(i,j) : mat.index(j,i); } 
        
        matrixScaledTransOp(const matrix& mat, double scale=1.0, CBLAS_TRANSPOSE op=CblasNoTrans)
        : mat(mat), scale(scale), op(op) {}
        
        matrixScaledTransOp(const scaled<matrix>& smat, CBLAS_TRANSPOSE op=CblasNoTrans)
        : mat(smat.data), scale(smat.scale), op(op) {}
        
        operator matrix() const; 
        
        //Scaling:
        matrixScaledTransOp& operator*=(double s) { scale *= s; return *this; }
        matrixScaledTransOp operator*(double s) const { return matrixScaledTransOp(mat,scale*s,op); }
        friend matrixScaledTransOp operator*(double s, const matrixScaledTransOp& A) { return A * s; }
};
matrix conj(const scaled<matrix>& A); 
matrixScaledTransOp dagger(const scaled<matrix>& A); 
matrixScaledTransOp transpose(const scaled<matrix>& A); 
matrix dagger_symmetrize(const scaled<matrix>& A); 
matrix transpose_symmetrize(const scaled<matrix>& A); 


//------------ Arithmetic ------------------

inline matrix& operator*=(matrix& m, double s) { scale(s, m); return m; }
inline scaled<matrix> operator*(double s, const matrix &m) { return scaled<matrix>(m, s); }
inline scaled<matrix> operator*(const matrix &m, double s) { return scaled<matrix>(m, s); }
inline scaled<matrix> operator-(const matrix &m) { return scaled<matrix>(m, -1); }
inline matrix& operator*=(matrix& m, complex s) { scale(s, m); return m; }
inline matrix operator*(complex s, const matrix &m) { matrix sm(m); sm *= s; return sm; }
inline matrix operator*(const matrix &m, complex s) { matrix sm(m); sm *= s; return sm; }
inline diagMatrix operator*(double s, const diagMatrix &m) { diagMatrix ret(m); return ret*=s; }
inline diagMatrix operator*(const diagMatrix &m, double s) { diagMatrix ret(m); return ret*=s; }
inline diagMatrix operator-(const diagMatrix &m) { return m * (-1); }

matrix operator*(const matrixScaledTransOp&, const matrixScaledTransOp&);
matrix operator*(const matrix&, const diagMatrix&);
matrix operator*(const diagMatrix&, const matrix&);
diagMatrix operator*(const diagMatrix&, const diagMatrix&);

inline matrix& operator+=(matrix &m1, const matrix &m2) { if(m1) axpy(1.0, m2, m1); else m1 = m2; return m1; }
inline matrix& operator-=(matrix &m1, const matrix &m2) { if(m1) axpy(-1.0, m2, m1); else m1 = -m2; return m1; }
inline matrix operator+(const matrix &m1, const matrix &m2) { matrix tm(m1); tm += m2; return tm; }
inline matrix operator-(const matrix &m1, const matrix &m2) { matrix tm(m1); tm -= m2; return tm; }
void axpy(double alpha, const diagMatrix& x, matrix& y);
inline matrix& operator+=(matrix& m, const diagMatrix& d) { if(m) axpy(1., d, m); else m = d; return m; }
inline matrix& operator-=(matrix& m, const diagMatrix& d) { if(m) axpy(-1., d, m); else m=-d; return m; }
inline matrix operator+(const matrix& m, const diagMatrix& d) { matrix ret(m); ret+=d; return ret; }
inline matrix operator+(const diagMatrix& d, const matrix& m) { matrix ret(m); ret+=d; return ret; }
inline matrix operator-(const matrix& m, const diagMatrix& d) { matrix ret(m); ret-=d; return ret; }
inline matrix operator-(const diagMatrix& d, const matrix& m) { matrix ret(-m); ret+=d; return ret; }
void axpy(double alpha, const diagMatrix& x, diagMatrix& y);
inline diagMatrix& operator+=(diagMatrix& d1, const diagMatrix& d2) { axpy(1., d2, d1); return d1; }
inline diagMatrix& operator-=(diagMatrix& d1, const diagMatrix& d2) { axpy(-1., d2, d1); return d1; }
inline diagMatrix operator+(const diagMatrix& d1, const diagMatrix& d2) { diagMatrix ret(d1); return ret+=d2; }
inline diagMatrix operator-(const diagMatrix& d1, const diagMatrix& d2) { diagMatrix ret(d1); return ret-=d2; }

//Functions required for using the Minimizable template on matrices:
diagMatrix clone(const diagMatrix& x);
matrix clone(const matrix& x);
double dot(const diagMatrix& x, const diagMatrix& y);
double dot(const matrix& x, const matrix& y);
void randomize(diagMatrix& x);
void randomize(matrix& x);


//------- Nonlinear matrix functions and their gradients ---------

matrix inv(const matrix& A);
diagMatrix inv(const diagMatrix& A);
matrix invApply(const matrix& A, const matrix& b); //return inv(A) * b (A must be hermitian, positive-definite)

matrix LU(const matrix& A);

complex det(const matrix& A);

double det(const diagMatrix& A);

matrix pow(const matrix& A, double exponent, matrix* Aevecs=0, diagMatrix* Aeigs=0);

matrix invsqrt(const matrix& A, matrix* Aevecs=0, diagMatrix* Aeigs=0);

matrix cis(const matrix& A, matrix* Aevecs=0, diagMatrix* Aeigs=0);

matrix cisInv(const matrix& A, matrix* Bevecs=0, diagMatrix* Beigs=0);

matrix sqrt_grad(const matrix& grad_sqrtA, const matrix& Aevecs, const diagMatrix& Aeigs);

matrix cis_grad(const matrix& grad_cisA, const matrix& Aevecs, const diagMatrix& Aeigs);


//------------ Misc matrix functions --------------
complex trace(const matrix &m);
double trace(const diagMatrix& m);
double nrm2(const diagMatrix& m);
diagMatrix diag(const matrix &m); 
diagMatrix eye(int N); 
matrix zeroes(int nRows, int nCols); 

class tiledBlockMatrix
{
        const matrix& mBlock; 
        int nBlocks; 
        const std::vector<complex>* phaseArr; 
public:
        tiledBlockMatrix(const matrix& mBlock, int nBlocks, const std::vector<complex>* phaseArr=0);
        
        matrix operator*(const matrix&) const; 
};

#endif  // JDFTX_ELECTRONIC_MATRIX_H