Line data Source code
1 : /* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
2 : Copyright (c) 2026 Nils E. Strand, Siyao Yang, Yuehaw Khoo, Aaron R. Dinner
3 :
4 : See the COPYRIGHT file distributed with this software for license details.
5 : +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
6 : #include "TTHelper.h"
7 : #include "tools/OpenMP.h"
8 : #include <cstdio>
9 : #include <stdexcept>
10 : #ifdef __PLUMED_HAS_LIBITENSOR
11 : #ifdef __PLUMED_HAS_LIBHDF5
12 :
13 : using namespace itensor;
14 :
15 : namespace PLMD {
16 : namespace ttsketch {
17 :
18 43 : void ttWrite(const std::string& filename, const MPS& tt, unsigned count) {
19 : // count starts at 2 on the first sketch call; open for writing on first call, append thereafter
20 43 : auto f = count == 2 ? h5_open(filename, 'w') : h5_open(filename, 'a');
21 172 : h5_write(f, "tt_" + std::to_string(count - 1), tt);
22 43 : }
23 :
24 8 : MPS ttRead(const std::string& filename, unsigned count) {
25 : auto f = h5_open(filename, 'r');
26 32 : return h5_read<MPS>(f, "tt_" + std::to_string(count - 1));
27 : }
28 :
29 0 : void ttSumWrite(const std::string& filename, const MPS& tt, unsigned count) {
30 : auto f = h5_open(filename, 'a');
31 0 : h5_write(f, "vb_" + std::to_string(count - 1), tt);
32 0 : }
33 :
34 0 : MPS ttSumRead(const std::string& filename, unsigned count) {
35 : auto f = h5_open(filename, 'r');
36 0 : return h5_read<MPS>(f, "vb_" + std::to_string(count - 1));
37 : }
38 :
39 5 : unsigned ttCount(const std::string& filename) {
40 : {
41 : // h5::file throws (and HDF5 prints an error stack) on a missing file
42 5 : FILE* fp = std::fopen(filename.c_str(), "r");
43 5 : if(!fp) {
44 : return 0;
45 : }
46 4 : std::fclose(fp);
47 : }
48 : auto f = h5_open(filename, 'r');
49 4 : h5::group g(f);
50 : unsigned n = 0;
51 18 : while(g.has_subgroup("tt_" + std::to_string(n + 1))) {
52 : ++n;
53 : }
54 : return n;
55 : }
56 :
57 13 : void ttMomentsWrite(const std::string& filename, const MPS& B, const std::vector<ITensor>& A, unsigned count) {
58 : auto f = h5_open(filename, 'a');
59 13 : const std::string tag = std::to_string(count - 1);
60 39 : h5_write(f, "bprev_" + tag, B);
61 31 : for(unsigned k = 0; k < A.size(); ++k) {
62 90 : h5_write(f, "aprev_" + tag + "_" + std::to_string(k + 1), A[k]);
63 : }
64 13 : }
65 :
66 3 : bool ttMomentsRead(const std::string& filename, unsigned count, unsigned nA, MPS& B, std::vector<ITensor>& A) {
67 : auto f = h5_open(filename, 'r');
68 3 : h5::group g(f);
69 3 : const std::string tag = std::to_string(count - 1);
70 6 : if(!g.has_subgroup("bprev_" + tag)) {
71 : return false;
72 : }
73 5 : for(unsigned k = 0; k < nA; ++k) {
74 9 : if(!g.has_subgroup("aprev_" + tag + "_" + std::to_string(k + 1))) {
75 : return false;
76 : }
77 : }
78 6 : B = h5_read<MPS>(f, "bprev_" + tag);
79 2 : A.resize(nA);
80 5 : for(unsigned k = 0; k < nA; ++k) {
81 15 : A[k] = h5_read<ITensor>(f, "aprev_" + tag + "_" + std::to_string(k + 1));
82 : }
83 : return true;
84 : }
85 :
86 : // Evaluate TT and gradient together using a prefix-suffix sweep: O(d) contractions.
87 : //
88 : // For a TT of length d:
89 : // L[i] = (G_1 * phi_1) * ... * (G_i * phi_i) [left partial product up to i]
90 : // R[i] = (G_i * phi_i) * ... * (G_d * phi_d) [right partial product from i]
91 : //
92 : // val = elt(L[d])
93 : // grad[k] = L[k-1] * (G_k * dphi_k) * R[k+1]
94 : std::pair<double, std::vector<double>>
95 60 : ttEvalAndGrad(const itensor::MPS& tt,
96 : const std::vector<BasisFunc>& basis,
97 : const std::vector<double>& elements,
98 : bool conv) {
99 60 : int d = length(tt);
100 60 : auto s = siteInds(tt);
101 60 : const unsigned nt = OpenMP::getNumThreads();
102 :
103 : // Cache cores and build phi/dphi — all independent across dimensions
104 60 : std::vector<ITensor> cores(d), phi(d), dphi(d);
105 60 : #pragma omp parallel for num_threads(nt) schedule(static)
106 : for (int i = 0; i < d; ++i) {
107 : cores[i] = tt(i + 1);
108 : phi[i] = dphi[i] = ITensor(s(i + 1));
109 : for (int j = 1; j <= dim(s(i + 1)); ++j) {
110 : phi[i].set(s(i + 1) = j, basis[i](elements[i], j, conv));
111 : dphi[i].set(s(i + 1) = j, basis[i].grad(elements[i], j, conv));
112 : }
113 : }
114 :
115 : // Prefix-suffix sweep
116 60 : std::vector<ITensor> L(d + 1), R(d + 1);
117 60 : L[0] = ITensor(1.0);
118 240 : for (int i = 0; i < d; ++i) {
119 180 : L[i + 1] = L[i] * (cores[i] * phi[i]);
120 : }
121 60 : R[d] = ITensor(1.0);
122 240 : for (int i = d - 1; i >= 0; --i) {
123 180 : R[i] = (cores[i] * phi[i]) * R[i + 1];
124 : }
125 :
126 60 : double val = elt(L[d]);
127 :
128 : // Gradient combination — independent across k
129 60 : std::vector<double> grad(d);
130 60 : #pragma omp parallel for num_threads(nt) schedule(static)
131 : for (int k = 0; k < d; ++k) {
132 : grad[k] = elt(L[k] * (cores[k] * dphi[k]) * R[k + 1]);
133 : }
134 :
135 60 : return {val, grad};
136 60 : }
137 :
138 96 : DenseTT::DenseTT(const MPS& tt) {
139 96 : const int d = length(tt);
140 96 : rank_.assign(d + 1, 1);
141 96 : nbasis_.resize(d);
142 96 : core_.resize(d);
143 333 : for(int i = 1; i <= d; ++i) {
144 237 : const Index s = siteIndex(tt, i);
145 237 : const Index l = i > 1 ? linkIndex(tt, i - 1) : Index();
146 237 : const Index r = i < d ? linkIndex(tt, i) : Index();
147 237 : const int rl = i > 1 ? dim(l) : 1;
148 237 : const int rr = i < d ? dim(r) : 1;
149 237 : const int n = dim(s);
150 237 : rank_[i] = rr;
151 237 : nbasis_[i - 1] = n;
152 : // the core has exactly the indices (l,s,r) without the missing boundary bonds
153 333 : if(order(tt(i)) != 1 + (i > 1) + (i < d)) {
154 0 : throw std::runtime_error("unexpected index structure in TT core " + std::to_string(i));
155 : }
156 : auto& c = core_[i - 1];
157 237 : c.resize(static_cast<size_t>(rl) * n * rr);
158 2501 : for(int a = 1; a <= rl; ++a) {
159 18070 : for(int j = 1; j <= n; ++j) {
160 372716 : for(int b = 1; b <= rr; ++b) {
161 : double v;
162 356910 : if(i > 1 && i < d) {
163 681308 : v = tt(i).elt(l = a, s = j, r = b);
164 16256 : } else if(i > 1) {
165 16270 : v = tt(i).elt(l = a, s = j);
166 8121 : } else if(i < d) {
167 16242 : v = tt(i).elt(s = j, r = b);
168 : } else {
169 0 : v = tt(i).elt(s = j);
170 : }
171 356910 : c[(static_cast<size_t>(a - 1) * n + (j - 1)) * rr + (b - 1)] = v;
172 : }
173 : }
174 : }
175 : }
176 96 : }
177 :
178 823 : double DenseTT::eval(const std::vector<BasisFunc>& basis, const std::vector<double>& elements, bool conv, std::vector<double>* grad) const {
179 823 : const unsigned d = core_.size();
180 : // contracted cores: M[i] = sum_j core_i[:,j,:] phi_j(x_i), D[i] the same with dphi
181 823 : thread_local std::vector<std::vector<double>> M, D;
182 823 : thread_local std::vector<double> phi, dphi;
183 823 : M.resize(d);
184 823 : D.resize(d);
185 2786 : for(unsigned i = 0; i < d; ++i) {
186 1963 : const int rl = rank_[i], rr = rank_[i + 1], n = nbasis_[i];
187 1963 : phi.resize(n);
188 1963 : dphi.resize(n);
189 1963 : basis[i].evalAll(elements[i], conv, phi.data(), dphi.data());
190 1963 : M[i].assign(static_cast<size_t>(rl) * rr, 0.0);
191 1963 : if(grad) {
192 696 : D[i].assign(static_cast<size_t>(rl) * rr, 0.0);
193 : }
194 : const double* c = core_[i].data();
195 14839 : for(int a = 0; a < rl; ++a) {
196 12876 : double* m = M[i].data() + static_cast<size_t>(a) * rr;
197 12876 : double* g = grad ? D[i].data() + static_cast<size_t>(a) * rr : nullptr;
198 101122 : for(int j = 0; j < n; ++j) {
199 88246 : const double* cj = c + (static_cast<size_t>(a) * n + j) * rr;
200 88246 : const double p = phi[j];
201 1839305 : for(int b = 0; b < rr; ++b) {
202 1751059 : m[b] += cj[b] * p;
203 : }
204 88246 : if(g) {
205 29029 : const double dp = dphi[j];
206 569502 : for(int b = 0; b < rr; ++b) {
207 540473 : g[b] += cj[b] * dp;
208 : }
209 : }
210 : }
211 : }
212 : }
213 : // left sweep: left[i] = M[0]...M[i-1] (row vector of size r_i)
214 823 : thread_local std::vector<std::vector<double>> left, right;
215 823 : left.resize(d + 1);
216 823 : left[0].assign(1, 1.0);
217 2786 : for(unsigned i = 0; i < d; ++i) {
218 1963 : const int rl = rank_[i], rr = rank_[i + 1];
219 1963 : left[i + 1].assign(rr, 0.0);
220 14839 : for(int a = 0; a < rl; ++a) {
221 12876 : const double la = left[i][a];
222 12876 : const double* m = M[i].data() + static_cast<size_t>(a) * rr;
223 263359 : for(int b = 0; b < rr; ++b) {
224 250483 : left[i + 1][b] += la * m[b];
225 : }
226 : }
227 : }
228 823 : const double val = left[d][0];
229 823 : if(grad) {
230 : // right sweep: right[i] = M[i]...M[d-1] (column vector of size r_i)
231 287 : right.resize(d + 1);
232 287 : right[d].assign(1, 1.0);
233 983 : for(int i = d - 1; i >= 0; --i) {
234 696 : const int rl = rank_[i], rr = rank_[i + 1];
235 696 : right[i].assign(rl, 0.0);
236 4904 : for(int a = 0; a < rl; ++a) {
237 4208 : const double* m = M[i].data() + static_cast<size_t>(a) * rr;
238 : double sum = 0.0;
239 81486 : for(int b = 0; b < rr; ++b) {
240 77278 : sum += m[b] * right[i + 1][b];
241 : }
242 4208 : right[i][a] = sum;
243 : }
244 : }
245 : // grad_i = left[i] * D[i] * right[i+1]
246 287 : grad->resize(d);
247 983 : for(unsigned i = 0; i < d; ++i) {
248 696 : const int rl = rank_[i], rr = rank_[i + 1];
249 : double sum = 0.0;
250 4904 : for(int a = 0; a < rl; ++a) {
251 4208 : const double* g = D[i].data() + static_cast<size_t>(a) * rr;
252 : double t = 0.0;
253 81486 : for(int b = 0; b < rr; ++b) {
254 77278 : t += g[b] * right[i + 1][b];
255 : }
256 4208 : sum += left[i][a] * t;
257 : }
258 696 : (*grad)[i] = sum;
259 : }
260 : }
261 823 : return val;
262 : }
263 :
264 : // ttEval delegates to ttEvalAndGrad and discards the gradient.
265 : // Called infrequently (error measurement, well-tempered height) so the
266 : // cost of computing dphi is negligible relative to avoiding code duplication.
267 0 : double ttEval(const MPS& tt,
268 : const std::vector<BasisFunc>& basis,
269 : const std::vector<double>& elements,
270 : bool conv) {
271 0 : return ttEvalAndGrad(tt, basis, elements, conv).first;
272 : }
273 :
274 : // Compute covariance matrix, marginal means, and partition function of the TT distribution.
275 : // Precomputes ITensors for the three moment integrals int0/int1/int2 per dimension,
276 : // then evaluates expectations as TT contractions with these integral vectors.
277 0 : std::tuple<Matrix<double>, std::vector<double>, double> covMat(const MPS& tt, const std::vector<BasisFunc>& basis) {
278 0 : int d = length(tt);
279 0 : auto s = siteInds(tt);
280 : // integral vectors: int0[i][j] = int phi_j(x_i) dx_i, etc.
281 0 : std::vector<ITensor> basis_int0(d), basis_int1(d), basis_int2(d);
282 0 : for(int i = 1; i <= d; ++i) {
283 0 : basis_int0[i - 1] = basis_int1[i - 1] = basis_int2[i - 1] = ITensor(s(i));
284 0 : for(int j = 1; j <= dim(s(i)); ++j) {
285 0 : basis_int0[i - 1].set(s(i) = j, basis[i - 1].int0(j));
286 0 : basis_int1[i - 1].set(s(i) = j, basis[i - 1].int1(j));
287 0 : basis_int2[i - 1].set(s(i) = j, basis[i - 1].int2(j));
288 : }
289 : }
290 : // Z = int tt(x) dx (partition function); normalize to get probability density rho
291 0 : auto Z = tt(1) * basis_int0[0];
292 0 : for(int i = 2; i <= d; ++i) {
293 0 : Z *= tt(i) * basis_int0[i - 1];
294 : }
295 0 : auto rho = tt;
296 0 : rho /= elt(Z);
297 : // ei[k] = E[x_k], eii[k] = E[x_k^2], eij[k][l] = E[x_k * x_l] for k < l
298 0 : std::vector<double> ei(d), eii(d);
299 0 : Matrix<double> eij(d, d);
300 0 : for(int k = 1; k <= d; ++k) {
301 : // replace dimension k's integral vector with int1 (resp. int2) to get E[x_k] (resp. E[x_k^2])
302 0 : auto eival = rho(1) * (k == 1 ? basis_int1[0] : basis_int0[0]);
303 0 : auto eiival = rho(1) * (k == 1 ? basis_int2[0] : basis_int0[0]);
304 0 : for(int i = 2; i <= d; ++i) {
305 0 : eival *= rho(i) * (k == i ? basis_int1[i - 1] : basis_int0[i - 1]);
306 0 : eiival *= rho(i) * (k == i ? basis_int2[i - 1] : basis_int0[i - 1]);
307 : }
308 0 : ei[k - 1] = elt(eival);
309 0 : eii[k - 1] = elt(eiival);
310 0 : for(int l = k + 1; l <= d; ++l) {
311 : // replace both dimensions k and l with int1 to get E[x_k * x_l]
312 0 : auto eijval = rho(1) * (k == 1 ? basis_int1[0] : basis_int0[0]);
313 0 : for(int i = 2; i <= d; ++i) {
314 0 : eijval *= rho(i) * (k == i || l == i ? basis_int1[i - 1] : basis_int0[i - 1]);
315 : }
316 0 : eij(k - 1, l - 1) = elt(eijval);
317 0 : }
318 0 : }
319 : // sigma(k,l) = Cov(x_k, x_l) = E[x_k*x_l] - E[x_k]*E[x_l]
320 0 : Matrix<double> sigma(d, d);
321 0 : for(int k = 1; k <= d; ++k) {
322 0 : for(int l = k; l <= d; ++l) {
323 0 : sigma(k - 1, l - 1) = sigma(l - 1, k - 1) = k == l ? eii[k - 1] - pow(ei[k - 1], 2) : eij(k - 1, l - 1) - ei[k - 1] * ei[l - 1];
324 : }
325 : }
326 0 : return std::make_tuple(sigma, ei, elt(Z));
327 0 : }
328 :
329 : // Compute the 2D marginal density of the normalized TT distribution for dimensions
330 : // pos1 and pos2 on a (bins x bins) grid. All other dimensions are integrated out
331 : // analytically using int0, which collapses those TT cores to scalar factors.
332 0 : void marginal2d(const MPS& tt, const std::vector<BasisFunc>& basis, int pos1, int pos2, Matrix<double>& grid, bool conv) {
333 0 : int bins = grid.nrows();
334 0 : int d = length(tt);
335 0 : auto s = siteInds(tt);
336 0 : std::vector<ITensor> basis_int0(d);
337 0 : for(int i = 1; i <= d; ++i) {
338 0 : basis_int0[i - 1] = ITensor(s(i));
339 0 : for(int j = 1; j <= dim(s(i)); ++j) {
340 0 : basis_int0[i - 1].set(s(i) = j, basis[i - 1].int0(j));
341 : }
342 : }
343 : // normalize to probability density
344 0 : auto Z = tt(1) * basis_int0[0];
345 0 : for(int i = 2; i <= d; ++i) {
346 0 : Z *= tt(i) * basis_int0[i - 1];
347 : }
348 0 : auto rho = tt;
349 0 : rho /= elt(Z);
350 0 : for(int i = 0; i < bins; ++i) {
351 0 : for(int j = 0; j < bins; ++j) {
352 0 : double x = basis[pos1 - 1].dom().first + i * (basis[pos1 - 1].dom().second - basis[pos1 - 1].dom().first) / bins;
353 0 : double y = basis[pos2 - 1].dom().first + j * (basis[pos2 - 1].dom().second - basis[pos2 - 1].dom().first) / bins;
354 0 : ITensor xevals(s(pos1)), yevals(s(pos2));
355 0 : for(int k = 1; k <= dim(s(pos1)); ++k) {
356 0 : xevals.set(s(pos1) = k, basis[pos1 - 1](x, k, conv));
357 0 : yevals.set(s(pos2) = k, basis[pos2 - 1](y, k, conv));
358 : }
359 0 : auto val = rho(1) * (pos1 == 1 ? xevals : basis_int0[0]);
360 0 : for(int k = 2; k <= d; ++k) {
361 0 : if(pos1 == k) {
362 0 : val *= rho(k) * xevals;
363 0 : } else if(pos2 == k) {
364 0 : val *= rho(k) * yevals;
365 : } else {
366 0 : val *= rho(k) * basis_int0[k - 1];
367 : }
368 : }
369 0 : grid(i, j) = elt(val);
370 0 : }
371 : }
372 0 : }
373 :
374 : }
375 : }
376 : #endif // __PLUMED_HAS_LIBHDF5
377 : #endif // __PLUMED_HAS_LIBITENSOR
|
