ColumnBundle.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_COLUMNBUNDLE_H
#define JDFTX_ELECTRONIC_COLUMNBUNDLE_H

#include <electronic/common.h>
#include <core/ManagedMemory.h>
#include <electronic/Basis.h>
#include <electronic/matrix.h>
#include <core/ScalarField.h>
#include <core/scaled.h>

class ColumnBundle : public ManagedMemory
{
        int ncols;
        size_t col_length; //typically equal to either basis->nbasis (or 2*basis->nbasis for spinors in noncollinear mode)

public:
        int nCols() const { return ncols; } 
        size_t colLength() const { return col_length; } 
        explicit operator bool() const { return ncols*col_length; } 

        size_t index(int i, size_t j) const { return i*col_length+j; }

        bool isSpinor() const { return basis && (col_length==2*basis->nbasis); }
        int spinorLength() const { return isSpinor() ? 2 : 1; }
        
        const QuantumNumber *qnum;
        const Basis *basis;

        void init(int nc, size_t len, const Basis* b, const QuantumNumber* q, bool onGpu=false); 
        void free(); 
        ColumnBundle(int nc=0, size_t len=0, const Basis* b=NULL, const QuantumNumber* q=NULL, bool onGpu=false); 
        ColumnBundle(const ColumnBundle&); 
        ColumnBundle(ColumnBundle&&); 

        ColumnBundle similar(int ncOverride=-1) const;
        
        ColumnBundle& operator=(const ColumnBundle&); 
        ColumnBundle& operator=(ColumnBundle&&); 
        
        // Get/set columns
        ColumnBundle getSub(int colStart, int colStop) const; 
        void setSub(int colStart, const ColumnBundle&); 
        
        complexScalarFieldTilde getColumn(int i, int s) const; 
        void setColumn(int i, int s, const complexScalarFieldTilde&); 
        void accumColumn(int i, int s, const complexScalarFieldTilde&); 
        
        void randomize(int colStart, int colStop); 
};

void init(std::vector<ColumnBundle>&, int nbundles, int ncols=0, const Basis* basis=0, const ElecInfo* eInfo=0);

void randomize(std::vector<ColumnBundle>&, const ElecInfo& eInfo);
void write(const std::vector<ColumnBundle>&, const char *fname, const ElecInfo& eInfo);

struct ColumnBundleReadConversion
{       bool realSpace; 
        int nBandsOld; 
        double Ecut, EcutOld; 
        vector3<int> S_old; 
        
        ColumnBundleReadConversion();
};
void read(std::vector<ColumnBundle>&, const char *fname, const ElecInfo& eInfo, const ColumnBundleReadConversion* conversion=0);

// Used in the CG template Minimize.h
ColumnBundle clone(const ColumnBundle&);  
void randomize(ColumnBundle& x); 
double dot(const ColumnBundle& x, const ColumnBundle& y); 


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

ColumnBundle& operator+=(ColumnBundle& Y, const scaled<ColumnBundle> &X);
ColumnBundle& operator-=(ColumnBundle& Y, const scaled<ColumnBundle> &X);
ColumnBundle operator+(const scaled<ColumnBundle> &Y1, const scaled<ColumnBundle> &Y2);
ColumnBundle operator-(const scaled<ColumnBundle> &Y1, const scaled<ColumnBundle> &Y2);

ColumnBundle& operator*=(ColumnBundle& X, double s);
ColumnBundle operator*(double s, ColumnBundle&& Y);
ColumnBundle operator*(ColumnBundle&& Y, double s);
scaled<ColumnBundle> operator*(double s, const ColumnBundle &Y);
scaled<ColumnBundle> operator*(const ColumnBundle &Y, double s);
scaled<ColumnBundle> operator-(const ColumnBundle &Y);
ColumnBundle& operator*=(ColumnBundle& X, complex s);
ColumnBundle operator*(complex s, const ColumnBundle &Y);
ColumnBundle operator*(const ColumnBundle &Y, complex s);

struct ColumnBundleMatrixProduct
{       const ColumnBundle& Y; 
        const matrixScaledTransOp& Mst; 
        double scale; 
        
        int nCols() const { return Mst.nCols(); } 
        size_t colLength() const { return Y.colLength(); } 

        //Scaling:
        ColumnBundleMatrixProduct& operator*=(double s) { scale *= s; return *this; }
        ColumnBundleMatrixProduct operator*(double s) const { return ColumnBundleMatrixProduct(Y,Mst,scale*s); }
        friend ColumnBundleMatrixProduct operator*(double s, const ColumnBundleMatrixProduct& A) { return A * s; }

        operator ColumnBundle() const; 
        void scaleAccumulate(double alpha, double beta, ColumnBundle& YM) const; 
private:
        ColumnBundleMatrixProduct(const ColumnBundle& Y, const matrixScaledTransOp& Mst, double scale=1.) : Y(Y), Mst(Mst), scale(scale) {}
        friend ColumnBundleMatrixProduct operator*(const scaled<ColumnBundle>& sY, const matrixScaledTransOp& Mst);
};

//Delay ColumnBundle * matrix and combine it with ColumnBundle accumulate operations when possible:
ColumnBundleMatrixProduct operator*(const scaled<ColumnBundle>& sY, const matrixScaledTransOp& Mst);
ColumnBundle& operator+=(ColumnBundle& Y, const ColumnBundleMatrixProduct &XM);
ColumnBundle& operator-=(ColumnBundle& Y, const ColumnBundleMatrixProduct &XM);
ColumnBundle operator+(const ColumnBundleMatrixProduct &XM1, const ColumnBundleMatrixProduct &XM2);
ColumnBundle operator-(const ColumnBundleMatrixProduct &XM1, const ColumnBundleMatrixProduct &XM2);

ColumnBundle operator*(const scaled<ColumnBundle>&, const diagMatrix&);
matrix operator^(const scaled<ColumnBundle>&, const scaled<ColumnBundle>&); 

#endif // JDFTX_ELECTRONIC_COLUMNBUNDLE_H