lr.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <generic/common.hpp>
#include <generic/storage.hpp>
#include <generic/kernels.hpp>
#include <generic/matrix.hpp>
 
namespace Ensign {
 
/* Class to store the low-rank factors (X, S, V) that are used in the computation.
*
* The low-rank factors can reside either on the host (CPU) or the device (GPU)
* depending on the value of stloc (either stloc::host or stloc::device).
*/
template<class T>
struct lr2 {
  multi_array<T, 2> S;
  multi_array<T, 2> X;
  multi_array<T, 2> V;
 
  lr2(stloc sl=stloc::host) : S(sl), X(sl), V(sl) {}
 
  lr2(Index r, array<Index,2> N, stloc sl=stloc::host) : S({r,r},sl), X({N[0],r},sl), V({N[1],r},sl) {}
 
  lr2(Index r, Index r_max, array<Index,2> N, stloc sl=stloc::host) : S({r,r},sl), X({N[0],r_max},{N[0],r},sl), V({N[1],r_max},{N[1],r},sl) {}
 
  void resize(Index r, array<Index,2> N) {
    S.resize({r,r});
    X.resize({N[0],r});
    V.resize({N[1],r});
  }
 
  void update_info(Index r) {
    X.update_shape({X.shape()[0],r});
    S.resize({r,r});
    V.update_shape({V.shape()[0],r});
  }
 
  Index size_X() const {
    return X.shape()[0];
  }
  Index size_V() const {
    return V.shape()[0];
  }
  Index rank() const {
    return S.shape()[0];
  }
 
  void set_zero() {
    S.set_zero();
  }
 
  /* Computes the full matrix (can be very expensive and should not be used in production code)
  */
  multi_array<T,2> full(const Ensign::Matrix::blas_ops& blas) const {
    multi_array<T, 2> K = X;
    blas.matmul(X, S, K);
    multi_array<T, 2> out({size_X(), size_V()});
    blas.matmul_transb(K, V, out);
    return out;
  }
};
 
/*  Add two low-rank representations together.
*
*   Note that no truncation is performed by this function. To do that call lr_truncate.
*/
template<class T, class IP>
void lr_add(T alpha, const lr2<T>& A, T beta, const lr2<T>& B, lr2<T>& out,
            IP inner_product_X,
            IP inner_product_V,
            const Ensign::Matrix::blas_ops& blas);
 
/*  Add three low-rank representations together.
*
*   Note that no truncation is performed by this function. To do that call lr_truncate.
*/
template<class T, class IP>
void lr_add(T alpha, const lr2<T>& A, T beta, const lr2<T>& B,
            T gamma, const lr2<T>& C, lr2<T>& out,
            IP inner_product_X,
            IP inner_product_V,
            const Ensign::Matrix::blas_ops& blas);
 
 
/* Multiplies two low-rank representations together.
*
*  Note that no truncation is performed by this function. To do that call lr_truncate.
*/
template<class T, class IP>
void lr_mul(const lr2<T>& A, const lr2<T>& B, lr2<T>& out,
            IP inner_product_X,
            IP inner_product_V,
            const Ensign::Matrix::blas_ops& blas);
 
/*  Truncate a low-rank representation inplace to a fixed rank r.
*
*   The rank is determined by the size of out.
*/
template<class T>
void lr_truncate(const lr2<T>& in, lr2<T>& out, const Ensign::Matrix::blas_ops& blas);
 
 
/* Computes the inner product of two low-rank representations.
*
*  Computes the inner product of two low-rank representations.
*/
template<class T>
double lr_inner_product(const lr2<T>& A, const lr2<T>&B, T w, const Ensign::Matrix::blas_ops& blas);
 
 
/* Computes the norm of a low-rank representations.
*
*  This assumes that X and V are orthogonalized. Thus, no inner product needs
*  to be passed to this function. The function simply computes
*  \f$\sum_{ij} S_{ij}^2\f$
*/
template<class T>
double lr_norm_sq(const lr2<T>& A, const Ensign::Matrix::blas_ops& blas);
 
 
 
 
/* Initializes (an already allocated) lr2 to \f$\sum_i X_i V_i\f$.
*
* The rank r is taken from the initialized lr. If X.size() and V.size() are smaller
* than r, the remaining basis functions are initialized in a random manner (subject
* (to the orthogonality condition).
* X.size() == V.size() is required.
*/
template<class T, class IP>
void initialize(lr2<T>& lr, vector<const T*> X, vector<const T*> V,
                IP inner_product_X,
                IP inner_product_V,
                const Ensign::Matrix::blas_ops& blas);
 
template<class T>
void initialize(lr2<T>& lr, vector<const T*> X, vector<const T*> V,
                array<double,3> h_xx, array<double,3> h_vv,
                const Ensign::Matrix::blas_ops& blas);
 
/* Return an inner product function object for use in, e.g., in gram_schmidt.
*/
template<class T>
std::function<T(T*,T*)> inner_product_from_const_weight(T w, Index N);
 
/* Return an inner product function object for use in, e.g., in gram_schmidt.
*
*/
template<class T>
std::function<T(T*,T*)> inner_product_from_weight(const T* w, Index N);
 
 
/* Orthogonalizes Q and returns the corresponding upper triangular matrix R.
*
* Note that the inputs are overwritten.
*/
 
struct orthogonalize {
 
  orthogonalize(const Ensign::Matrix::blas_ops* _blas);
  ~orthogonalize();
 
  void operator()(multi_array<double,2>& Q, multi_array<double,2>& R, std::function<double(double*,double*)> inner_product);
  void operator()(multi_array<double,2>& Q, multi_array<double,2>& R, double w);
  void operator()(multi_array<double,2>& Q, multi_array<double,2>& R, double* w);
  
  // Forces the use of Gram Schmidt.
  // This guarantees that if the first number of basis functions are already orthonormal, they are not changed.
  void gram_schmidt(multi_array<double,2>& Q, multi_array<double,2>& R, std::function<double(double*,double*)> inner_product);
  void gram_schmidt(multi_array<double,2>& Q, multi_array<double,2>& R, double w);
  
private:
  const Ensign::Matrix::blas_ops* blas;
  #ifdef __CUDA__
  curandGenerator_t gen;
  #endif
};
 
} // namespace Ensign