LinOp.hpp

Go to the documentation of this file.

#ifndef CYTNX_LINOP_H_
#define CYTNX_LINOP_H_
 
#include "Type.hpp"
#include "cytnx_error.hpp"
#include <vector>
#include <fstream>
#include <functional>
#include <map>
#include <utility>
#include <algorithm>
#include "intrusive_ptr_base.hpp"
#include "Tensor.hpp"
#include "UniTensor.hpp"
#include "utils/vec_clone.hpp"
 
#ifdef BACKEND_TORCH
#else
namespace cytnx {
 
  class LinOp {
   private:
    // 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) {
      if (type == "mv") {
      } else if (type == "mv_elem") {
      } 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);
 
    // this expose to interface:
    virtual UniTensor matvec(const UniTensor &Tin);
    // virtual std::vector<UniTensor> matvec(const std::vector<UniTensor> &Tin);
  };
 
}  // namespace cytnx
 
#endif  // BACKEND_TORCH
 
#endif  // CYTNX_LINOP_H_