LinOp.hpp

Go to the documentation of this file.

#ifndef _H_LinOp_
#define _H_LinOp_
 
#include "Type.hpp"
#include "cytnx_error.hpp"
#include "Tensor.hpp"
#include "Scalar.hpp"
#include <vector>
#include <fstream>
#include <functional>
#include <map>
#include <utility>
#include <algorithm>
#include "intrusive_ptr_base.hpp"
#include "utils/vec_clone.hpp"
 
namespace cytnx{
    
    class LinOp{
        private:
 
            // function pointer:
            std::function<Tensor(const Tensor&)> _mvfunc;
        
            // type:
            std::string _type;
 
            // nx
            cytnx_uint64 _nx;
            
            // device
            int _device;
            int _dtype;
 
            // pre-storage data:
            std::map<cytnx_uint64,std::pair<std::vector<cytnx_uint64>,Tensor> > _elems; //map[i] -> pair[<js>,<Storage>]
            std::map<cytnx_uint64,std::pair<std::vector<cytnx_uint64>,Tensor> >::iterator _elems_it; 
 
            Tensor _mv_elemfunc(const Tensor &);
 
 
        public:
 
        // we need driver of void f(nx,vin,vout)
 
        LinOp(const std::string &type, const cytnx_uint64 &nx, const int &dtype=Type.Double, const int &device=Device.cpu, std::function<Tensor(const Tensor&)> custom_f = nullptr){
 
            if(type=="mv"){
                this->_mvfunc=custom_f;
 
            }else if(type=="mv_elem"){
                cytnx_error_msg(custom_f != nullptr,"[ERROR][LinOp] with type=mv_elem cannot accept any function. use set_elem/set_elemts instead.%s","\n");
                    
            
            }else
                cytnx_error_msg(type!="mv","[ERROR][LinOp] currently only type=\"mv\" (matvec) can be used.%s","\n");
 
            this->_type = type; 
            this->_nx = nx;
            cytnx_error_msg(device<-1 || device >=Device.Ngpus,"[ERROR] invalid device.%s","\n");
            this->_device = device;
            cytnx_error_msg(dtype<1 || dtype >= N_Type,"[ERROR] invalid dtype.%s","\n");
            this->_dtype = dtype;
        };
        void set_func(std::function<Tensor(const Tensor&)> custom_f, const int &dtype, const int &device){
            if(this->_type=="mv"){
                this->_mvfunc = custom_f;
                cytnx_error_msg(device<-1 || device >=Device.Ngpus,"[ERROR] invalid device.%s","\n");
                this->_device = device;
                cytnx_error_msg(dtype<1 || dtype >= N_Type,"[ERROR] invalid dtype.%s","\n");
                this->_dtype = dtype;
            }else{
                cytnx_error_msg(true,"[ERROR] cannot specify func with type=mv_elem%s. use set_elem instead.","\n");
            }
        };
 
        template<class T>
        void set_elem(const cytnx_uint64 &i, const cytnx_uint64 &j, const T &elem, const bool check_exists=true){
            this->_elems_it = this->_elems.find(i);
            if(this->_elems_it == this->_elems.end()){
                //not exists:
                Tensor x({1},this->_dtype);
                x(0) = elem;                
                this->_elems[i] = std::pair<std::vector<cytnx_uint64>,Tensor>({j},x);
 
            }else{
                std::vector<cytnx_uint64> &vi = this->_elems_it->second.first; // pair:
                Tensor &ie = this->_elems_it->second.second;
                if(check_exists){
                    cytnx_error_msg(std::find(vi.begin(), vi.end(), j)!=vi.end(),"[ERROR] the element is set%s","\n");
                }
                vi.push_back(j);
                ie.append(elem);
            }
 
        };
        Tensor::Tproxy operator()(const cytnx_uint64 &i, const cytnx_uint64 &j){
            //[Note that this can only call by mv_elem]
            //if the element is not exists, it will create one.
            this->_elems_it = this->_elems.find(i);
            if(this->_elems_it == this->_elems.end()){
                //not exists:
                Tensor x({1},this->_dtype);
                x(0) = 0;                
                this->_elems[i] = std::pair<std::vector<cytnx_uint64>,Tensor>({j},x);
                return this->_elems[i].second(0);
            }else{
                std::vector<cytnx_uint64> &vi = this->_elems_it->second.first; // pair:
                Tensor &ie = this->_elems_it->second.second;
                auto tmp_it = std::find(vi.begin(), vi.end(), j);
                
                //if(check_exists){
                //    cytnx_error_msg(std::find(vi.begin(), vi.end(), j)!=vi.end(),"[ERROR] the element is set%s","\n");
                //}
                if(tmp_it==vi.end()){
                    vi.push_back(j);
                    ie.append(0);
                    return ie(vi.size()-1);
                }else{
                    return ie(std::distance(vi.begin(), tmp_it));
                }
            }
        }
 
        void set_device(const int &device){
            cytnx_error_msg(device<-1 || device >=Device.Ngpus,"[ERROR] invalid device.%s","\n");
            this->_device = device;
        };
        void set_dtype(const int &dtype){
            cytnx_error_msg(dtype<1 || dtype >= N_Type,"[ERROR] invalid dtype.%s","\n");
            this->_dtype = dtype;
        };
        int device() const{
            return this->_device;
        };
        int dtype() const{
            return this->_dtype;
        };
        cytnx_uint64 nx() const{
            return this->_nx;
        };
 
        void _print();
 
        // this expose to interitance:
        // need user to check the output to be Tensor
        virtual Tensor matvec(const Tensor &Tin); 
    
        
       
    };
 
}
 
 
#endif