Line data    Source code

       1             : #ifndef __PLUMED_ttsketch_TTHelper_h
       2             : #define __PLUMED_ttsketch_TTHelper_h
       3             : #include "BasisFunc.h"
       4             : #include "tools/Matrix.h"
       5             : #ifdef __PLUMED_HAS_LIBITENSOR
       6             : #ifdef __PLUMED_HAS_LIBHDF5
       7             : #include <itensor/all.h>
       8             : #include <utility>
       9             : 
      10             : namespace PLMD {
      11             : namespace ttsketch {
      12             : 
      13             : // Write TT to HDF5 file under dataset "tt_{count-1}".
      14             : // count=2 opens the file for writing (creating/truncating); subsequent calls append.
      15             : void ttWrite(const std::string& filename, const itensor::MPS& tt, unsigned count);
      16             : // Read TT from HDF5 file at dataset "tt_{count-1}".
      17             : itensor::MPS ttRead(const std::string& filename, unsigned count);
      18             : // Append the cumulative bias TT (vb) to HDF5 file under dataset "vb_{count-1}".
      19             : void ttSumWrite(const std::string& filename, const itensor::MPS& tt, unsigned count);
      20             : // Read cumulative bias TT from HDF5 file at dataset "vb_{count-1}".
      21             : itensor::MPS ttSumRead(const std::string& filename, unsigned count);
      22             : // Number of TTs stored in the file, i.e. the largest k such that "tt_1".."tt_k" all exist
      23             : // (0 if the file does not exist).
      24             : unsigned ttCount(const std::string& filename);
      25             : // Append the accumulated tensor moments of the NONINTRUSIVE variant (B and the cross-gram
      26             : // matrices A) under "bprev_{count-1}" and "aprev_{count-1}_{k}", k = 1..A.size().
      27             : void ttMomentsWrite(const std::string& filename, const itensor::MPS& B, const std::vector<itensor::ITensor>& A, unsigned count);
      28             : // Read what ttMomentsWrite() wrote; returns false if the moments are not in the file.
      29             : bool ttMomentsRead(const std::string& filename, unsigned count, unsigned nA, itensor::MPS& B, std::vector<itensor::ITensor>& A);
      30             : // Evaluate the TT bias at point `elements` by contracting each core with its
      31             : // 1D basis vector. If conv=true, basis functions are evaluated in convolution mode.
      32             : double ttEval(const itensor::MPS& tt, const std::vector<BasisFunc>& basis, const std::vector<double>& elements, bool conv);
      33             : // Compute the covariance matrix, mean vector, and partition function of the TT distribution.
      34             : // Normalizes tt by Z = int tt(x) dx, then returns:
      35             : //   sigma(k,l) = E[x_k * x_l] - E[x_k]*E[x_l]  (covariance matrix)
      36             : //   mu[k]      = E[x_k]                           (marginal means)
      37             : //   Z          = int tt(x) dx                     (partition function)
      38             : // All expectations use int0/int1/int2 evaluated analytically via BasisFunc.
      39             : std::tuple<Matrix<double>, std::vector<double>, double> covMat(const itensor::MPS& tt, const std::vector<BasisFunc>& basis);
      40             : // Fill `grid` (bins x bins) with the 2D marginal density of `tt` for dimensions
      41             : // pos1 and pos2, obtained by integrating out all other dimensions analytically.
      42             : void marginal2d(const itensor::MPS& tt, const std::vector<BasisFunc>& basis, int pos1, int pos2, Matrix<double>& grid, bool conv);
      43             : // Plain-array copy of a TT, used to evaluate the bias and its gradient at every step.
      44             : // Core i is stored as core_[i][(a*n_i + j)*r_i + b], with a/b the left/right bond
      45             : // indices (r_0 = r_d = 1) and j the basis index. Building it costs O(sum r n r) ITensor
      46             : // element accesses, so it is rebuilt only when the TT changes.
      47             : class DenseTT {
      48             :   std::vector<int> rank_;                 // bond dimensions r_0..r_d
      49             :   std::vector<int> nbasis_;               // site dimensions n_1..n_d
      50             :   std::vector<std::vector<double>> core_; // the cores
      51             : public:
      52           0 :   DenseTT() = default;
      53             :   explicit DenseTT(const itensor::MPS& tt);
      54             :   bool empty() const {
      55             :     return core_.empty();
      56             :   }
      57             :   unsigned size() const {
      58             :     return core_.size();
      59             :   }
      60             :   int nbasis(unsigned i) const {
      61         113 :     return nbasis_[i];
      62             :   }
      63             :   // bond dimension r_i, i = 0..d (r_0 = r_d = 1)
      64             :   int rank(unsigned i) const {
      65         150 :     return rank_[i];
      66             :   }
      67             :   // core i (0-based), element (a,j,b) at (a*nbasis(i) + j)*rank(i+1) + b
      68             :   const std::vector<double>& core(unsigned i) const {
      69             :     return core_[i];
      70             :   }
      71             :   // Value of the TT at `elements`; if grad is not null it is resized to d and filled
      72             :   // with the gradient. Same result as ttEvalAndGrad(), at O(d n r^2) cost.
      73             :   double eval(const std::vector<BasisFunc>& basis, const std::vector<double>& elements, bool conv, std::vector<double>* grad) const;
      74             : };
      75             : std::pair<double, std::vector<double>>
      76             :                                     ttEvalAndGrad(const itensor::MPS& tt,
      77             :                                         const std::vector<BasisFunc>& basis,
      78             :                                         const std::vector<double>& elements,
      79             :                                         bool conv);
      80             : }
      81             : }
      82             : #endif // __PLUMED_HAS_LIBHDF5
      83             : #endif // __PLUMED_HAS_LIBITENSOR
      84             : #endif