Line data Source code
1 : /* Copyright (c) 2026, Nils E. Strand, Siyao Yang, Yuehaw Khoo, and Aaron R. Dinner
2 : *
3 : * Tensor train (TT)-metadynamics.
4 : * https://arxiv.org/abs/2603.13549
5 : * See module.md for installation instructions.
6 : */
7 :
8 : #include "TTHelper.h"
9 : #include "bias/Bias.h"
10 : #include "core/ActionRegister.h"
11 : #include "core/ActionSet.h"
12 : #include "core/PlumedMain.h"
13 : #include "tools/Exception.h"
14 : #include "tools/Communicator.h"
15 : #include "tools/File.h"
16 : #include "tools/OpenMP.h"
17 : #include "tools/OFile.h"
18 : #include "tools/IFile.h"
19 : #include <algorithm>
20 : #include <ctime>
21 : #include <filesystem>
22 : #include <random>
23 :
24 : namespace PLMD {
25 : namespace ttsketch {
26 :
27 : //+PLUMEDOC BIAS TTMETAD
28 : /*
29 : Performs metadynamics with a tensor-train representation of the bias potential.
30 :
31 : This action runs metadynamics while periodically compressing the accumulated
32 : Gaussian hills into a low-rank tensor train (TT) using the TT-Sketch algorithm.
33 : The TT representation keeps memory and evaluation cost bounded throughout long
34 : simulations, in contrast to standard metadynamics where cost grows linearly with
35 : the number of deposited hills.
36 :
37 : The basis functions of each argument are defined on [INTERVAL_MIN, INTERVAL_MAX] and are
38 : periodic over this interval. For a periodic argument the interval should thus coincide with the
39 : periodicity domain. For a non periodic argument the interval should contain the region that is
40 : explored: outside it the bias is the one at the closest limit and there is no biasing force along
41 : that argument, and hills deposited outside are folded back into the interval when they are
42 : compressed into the TT (a warning is written in the log).
43 :
44 : After every compression the TT is appended to the HDF5 file given by TT_FILE, and the hills file
45 : restarts empty. When restarting, the last TT in TT_FILE is read back together with the hills in FILE,
46 : i.e. those deposited after it. With multiple replicas each replica has its own TT_FILE (with the
47 : replica suffix), unless WALKERS_MPI is used, in which case the walkers share one TT_FILE next to
48 : the shared hills file. Two TTMETAD actions must use different TT_FILE.
49 :
50 : ## Examples
51 :
52 : The following input runs TT-MetaD on two distances. Hills are deposited every
53 : 500 steps and compressed into a TT every 1e6 steps with target rank 10.
54 :
55 : ```plumed
56 : d1: DISTANCE ATOMS=1,2
57 : d2: DISTANCE ATOMS=3,4
58 : TTMETAD ...
59 : ARG=d1,d2
60 : SIGMA=0.1,0.1
61 : HEIGHT=1.0
62 : PACE=500
63 : SKETCH_PACE=1000000
64 : SKETCH_RANK=10
65 : SKETCH_INITRANK=50
66 : INTERVAL_MIN=0.0,0.0
67 : INTERVAL_MAX=3.0,3.0
68 : FILE=HILLS
69 : LABEL=tt
70 : ... TTMETAD
71 : PRINT ARG=tt.bias FILE=COLVAR STRIDE=100
72 : ```
73 :
74 : */
75 : //+ENDPLUMEDOC
76 :
77 : using namespace PLMD::bias;
78 :
79 : #if !defined(__PLUMED_HAS_LIBITENSOR) || !defined(__PLUMED_HAS_LIBHDF5)
80 :
81 : class TTMetaD : public Bias {
82 : public:
83 : static void registerKeywords(Keywords& keys);
84 : explicit TTMetaD(const ActionOptions& ao) : Action(ao), Bias(ao) {
85 : error("TTMETAD requires libitensor and libhdf5. "
86 : "Reconfigure with --enable-libitensor.");
87 : }
88 : void calculate() override {}
89 : void update() override {}
90 : };
91 :
92 : #else
93 :
94 : using namespace itensor;
95 :
96 : class TTMetaD : public Bias {
97 :
98 : private:
99 : // singular values of the cross-gram matrices smaller than this fraction of the largest
100 : // one are not inverted when the TT cores are recovered (see paraSketch())
101 : static constexpr double svdTolerance = 1e-12;
102 : struct Gaussian {
103 : double height;
104 : std::vector<double> center;
105 : std::vector<double> sigma;
106 : std::vector<double> invsigma;
107 355 : Gaussian(const double h, const std::vector<double>& c, const std::vector<double>& s)
108 355 : : height(h), center(c), sigma(s), invsigma(s) {
109 1135 : for(unsigned i = 0; i < invsigma.size(); ++i) {
110 780 : if(abs(invsigma[i]) > 1.e-20) {
111 780 : invsigma[i] = 1.0 / invsigma[i];
112 : } else {
113 0 : invsigma[i] = 0.0;
114 : }
115 : }
116 355 : }
117 : };
118 : // standard metadynamics parameters
119 : double kbt_;
120 : int stride_;
121 : bool welltemp_;
122 : double biasf_;
123 : std::string fmt_;
124 : bool isFirstStep_;
125 : double height0_;
126 : std::vector<double> sigma0_;
127 : std::vector<Gaussian> hills_; // Gaussians deposited since the last TT update
128 : OFile hillsOfile_;
129 : std::string mw_dir_;
130 : std::string hillsfname_;
131 : bool walkers_mpi_;
132 : int mpi_size_;
133 : int mpi_rank_;
134 : // TT-sketch parameters
135 : unsigned d_; // number of CV dimensions
136 : int sketch_rc_; // initial (sketch) rank for the random TT coefficient tensor
137 : int sketch_r_; // target TT bond dimension after SVD truncation (0 = use cutoff)
138 : double sketch_cutoff_; // relative SVD truncation threshold (used when sketch_r_=0)
139 : int sketch_stride_; // steps between TT bias updates (tau in the paper)
140 : double sketch_alpha_; // weight for non-constant basis functions in the sketch coefficient TT
141 : std::vector<BasisFunc> sketch_basis_; // one BasisFunc per CV dimension
142 : unsigned sketch_count_; // number of TT updates performed so far (starts at 1, increments after each update)
143 : MPS vb_; // current TT approximation of the accumulated bias (V_bias^TT)
144 : DenseTT vbDense_; // plain-array copy of vb_ used to evaluate the bias (rebuilt by setVb())
145 : std::string ttfilename_; // HDF5 file where the TTs are stored (TT_FILE, with replica suffix or in the walkers directory)
146 : std::vector<double> imin_; // INTERVAL_MIN
147 : std::vector<double> imax_; // INTERVAL_MAX
148 : double sketch_until_; // simulation time after which the bias is frozen (no further updates)
149 : bool frozen_; // true once the bias has been frozen (after sketch_until_)
150 : bool sketch_conv_; // true if Gaussian smoothing (conv mode) is active for basis evaluation
151 : bool nonintrusive_; // if true, uses the intrusive-free variant: accumulates Bemp directly
152 : // instead of re-projecting the previous vb_ through the Gram matrix
153 : bool deterministic_; // if true, use a fixed RNG seed (42) in createTTCoeff() for reproducibility
154 : MPS B_prev_; // accumulated tensor moment from all previous sketches (nonintrusive only)
155 : std::vector<ITensor> A_prev_; // accumulated cross-gram matrices per bond (nonintrusive only)
156 :
157 : void readGaussians(IFile *ifile);
158 : void setVb(const MPS& vb);
159 : // clamp cv to [INTERVAL_MIN, INTERVAL_MAX]; out[i] is true if cv[i] was outside
160 : void clampToInterval(std::vector<double>& cv, std::vector<char>& out) const;
161 : bool writerReplica() const {
162 123 : return !this->walkers_mpi_ || this->mpi_rank_ == 0;
163 : }
164 : void writeGaussian(const Gaussian& hill, OFile& file);
165 : double getHeight(const std::vector<double>& cv);
166 : double getBias(const std::vector<double>& cv);
167 : double getBiasAndDerivatives(const std::vector<double>& cv, std::vector<double>& der);
168 : double evaluateGaussian(const std::vector<double>& cv, const Gaussian& hill);
169 : double evaluateGaussianAndDerivatives(const std::vector<double>& cv, const Gaussian& hill, std::vector<double>& der, std::vector<double>& dp);
170 : bool scanOneHill(IFile* ifile, std::vector<Value>& tmpvalues, std::vector<double>& center, std::vector<double>& sigma, double& height);
171 : void paraSketch();
172 : MPS createTTCoeff() const;
173 : std::vector<Matrix<double>> intBasisSample() const;
174 : std::pair<MPS, std::vector<Matrix<double>>> formTensorMoment(const std::vector<Matrix<double>>& M, const MPS& coeff) const;
175 : std::tuple<MPS, std::vector<ITensor>, std::vector<ITensor>> formTensorMomentVb(const MPS& coeff);
176 :
177 : public:
178 : explicit TTMetaD(const ActionOptions&);
179 : void calculate() override;
180 : void update() override;
181 : static void registerKeywords(Keywords& keys);
182 : };
183 :
184 16 : TTMetaD::TTMetaD(const ActionOptions& ao):
185 : PLUMED_BIAS_INIT(ao),
186 16 : kbt_(0.0),
187 16 : stride_(0),
188 16 : welltemp_(false),
189 16 : biasf_(-1.0),
190 16 : isFirstStep_(true),
191 16 : height0_(std::numeric_limits<double>::max()),
192 16 : mw_dir_(""),
193 16 : walkers_mpi_(false),
194 16 : mpi_size_(0),
195 16 : mpi_rank_(0),
196 16 : sketch_r_(0),
197 16 : sketch_cutoff_(0.0),
198 16 : sketch_count_(1),
199 16 : sketch_until_(std::numeric_limits<double>::max()),
200 16 : frozen_(false),
201 16 : sketch_conv_(false),
202 16 : nonintrusive_(false),
203 32 : deterministic_(false) {
204 : bool kernel;
205 32 : parseFlag("KERNEL_BASIS", kernel);
206 16 : this->d_ = getNumberOfArguments();
207 16 : if(this->d_ < 2) {
208 0 : error("Number of arguments must be at least 2");
209 : }
210 16 : parse("FMT", this->fmt_);
211 32 : parseVector("SIGMA", this->sigma0_);
212 16 : if(this->sigma0_.size() != d_) {
213 0 : error("number of arguments does not match number of SIGMA parameters");
214 : }
215 16 : parse("HEIGHT", this->height0_);
216 16 : parse("PACE", this->stride_);
217 16 : if(stride_ <= 0) {
218 0 : error("frequency for hill addition is nonsensical");
219 : }
220 16 : this->hillsfname_ = "HILLS";
221 16 : parse("FILE", this->hillsfname_);
222 16 : parse("BIASFACTOR", this->biasf_);
223 16 : if(this->biasf_ < 1.0 && this->biasf_ != -1.0) {
224 0 : error("well tempered bias factor is nonsensical");
225 : }
226 16 : this->kbt_ = getkBT();
227 16 : if(this->biasf_ >= 1.0) {
228 2 : if(this->kbt_ == 0.0) {
229 0 : error("Unless the MD engine passes the temperature to plumed, with well-tempered metad you must specify it using TEMP");
230 : }
231 2 : this->welltemp_ = true;
232 : }
233 :
234 16 : parseFlag("WALKERS_MPI", this->walkers_mpi_);
235 16 : parse("WALKERS_DIR", this->mw_dir_);
236 16 : if(this->walkers_mpi_ && this->mw_dir_== "") {
237 4 : const std::string ret = std::filesystem::current_path();
238 4 : this->mw_dir_ = ret + "/";
239 : // multi_sim_comm is valid only on the first rank of each replica
240 4 : if(comm.Get_rank() == 0) {
241 2 : multi_sim_comm.Bcast(this->mw_dir_, 0);
242 : }
243 4 : comm.Bcast(this->mw_dir_, 0);
244 : }
245 16 : if(this->walkers_mpi_ && this->mw_dir_ != "") {
246 8 : this->hillsfname_ = this->mw_dir_ + "/" + this->hillsfname_;
247 : }
248 16 : std::string ttfile = "ttsketch.h5";
249 16 : parse("TT_FILE", ttfile);
250 16 : if(this->walkers_mpi_) {
251 : // shared by all the walkers, next to the shared hills file
252 8 : this->ttfilename_ = this->mw_dir_ + "/" + ttfile;
253 : } else {
254 : // one file per replica, as for the other output files
255 24 : this->ttfilename_ = FileBase::appendSuffix(ttfile, plumed.getSuffix());
256 : }
257 19 : for(const auto* other : plumed.getActionSet().select<const TTMetaD*>()) {
258 3 : if(other->ttfilename_ == this->ttfilename_) {
259 0 : error("TT_FILE " + this->ttfilename_ + " is already used by action " + other->getLabel() + ": use a different TT_FILE");
260 : }
261 : }
262 :
263 16 : parse("SKETCH_RANK", this->sketch_r_);
264 16 : parse("SKETCH_CUTOFF", this->sketch_cutoff_);
265 16 : if(this->sketch_r_ <= 0 && (this->sketch_cutoff_ <= 0.0 || this->sketch_cutoff_ > 1.0)) {
266 0 : error("Valid SKETCH_RANK or SKETCH_CUTOFF needs to be specified");
267 : }
268 16 : parse("SKETCH_INITRANK", this->sketch_rc_);
269 16 : if(this->sketch_rc_ <= 0) {
270 0 : error("SKETCH_INITRANK must be positive");
271 : }
272 16 : parse("SKETCH_PACE", this->sketch_stride_);
273 16 : if(this->sketch_stride_ <= 0) {
274 0 : error("SKETCH_PACE must be positive");
275 : }
276 16 : std::vector<double> interval_min;
277 32 : parseVector("INTERVAL_MIN", interval_min);
278 16 : if(interval_min.size() != this->d_) {
279 0 : error("Number of arguments does not match number of INTERVAL_MIN parameters");
280 : }
281 16 : std::vector<double> interval_max;
282 32 : parseVector("INTERVAL_MAX", interval_max);
283 16 : if(interval_max.size() != this->d_) {
284 0 : error("Number of arguments does not match number of INTERVAL_MAX parameters");
285 : }
286 16 : int nbasis = 20;
287 16 : parse("SKETCH_NBASIS", nbasis);
288 16 : if(nbasis <= 1) {
289 0 : error("SKETCH_NBASIS must be greater than 1");
290 : }
291 16 : if(!kernel && nbasis % 2 == 0) {
292 0 : ++nbasis;
293 : }
294 16 : parse("SKETCH_ALPHA", this->sketch_alpha_);
295 16 : if(this->sketch_alpha_ <= 0.0 || this->sketch_alpha_ > 1.0) {
296 0 : error("SKETCH_ALPHA must be positive and no greater than 1");
297 : }
298 16 : std::vector<double> w;
299 32 : parseVector("SKETCH_WIDTH", w);
300 16 : if(w.size() == 0) {
301 14 : w.assign(this->d_, 0.0);
302 : }
303 16 : if(w.size() != this->d_) {
304 0 : error("Number of arguments does not match number of SKETCH_WIDTH parameters");
305 : }
306 52 : for (double val : w) {
307 36 : if (val != 0.0) {
308 6 : this->sketch_conv_ = true;
309 : }
310 : }
311 16 : std::vector<double> dx;
312 32 : parseVector("KERNEL_DX", dx);
313 16 : if(dx.size() == 0) {
314 15 : dx.resize(this->d_, 0.0);
315 : }
316 16 : if(dx.size() != this->d_) {
317 0 : error("Number of arguments does not match number of KERNEL_DX parameters");
318 : }
319 52 : for(unsigned i = 0; i < this->d_; ++i) {
320 36 : if(this->sketch_conv_ && w[i] <= 0.0) {
321 0 : error("Gaussian smoothing requires positive WIDTH");
322 : }
323 36 : if(kernel && dx[i] < 0.0) {
324 0 : error("Kernel basis requires positive KERNEL_DX");
325 : }
326 36 : if(interval_max[i] <= interval_min[i]) {
327 0 : error("INTERVAL_MAX parameters need to be greater than respective INTERVAL_MIN parameters");
328 : }
329 36 : if(getPntrToArgument(i)->isPeriodic()) {
330 : // the basis is periodic over the interval: for a periodic CV the two periods should coincide
331 : double pmin, pmax;
332 24 : getPntrToArgument(i)->getDomain(pmin, pmax);
333 24 : const double tol = 1e-4 * (pmax - pmin);
334 24 : if(std::abs(interval_min[i] - pmin) > tol || std::abs(interval_max[i] - pmax) > tol) {
335 0 : warning("the interval of the periodic argument " + getPntrToArgument(i)->getName() +
336 : " does not coincide with its periodicity domain: the bias will be discontinuous");
337 : }
338 : }
339 72 : this->sketch_basis_.push_back(BasisFunc(std::make_pair(interval_min[i], interval_max[i]), nbasis, w[i], kernel, dx[i]));
340 : }
341 16 : if(kernel && this->sketch_conv_) {
342 0 : error("kernel smoothing incompatible with kernel basis");
343 : }
344 16 : this->imin_ = interval_min;
345 16 : this->imax_ = interval_max;
346 16 : if(kernel) {
347 4 : for(unsigned i = 0; i < this->d_; ++i) {
348 3 : const BasisFunc& b = this->sketch_basis_[i];
349 3 : log.printf(" kernel basis of argument %s: %d kernels of width %g, Gram matrix condition number %.1e, %d nearly dependent direction(s) left out\n",
350 : getPntrToArgument(i)->getName().c_str(), nbasis - 1, b.dx(), b.gramCondition(), b.gramDropped());
351 3 : if(b.gramDropped() > 1) {
352 0 : warning("the kernel basis of argument " + getPntrToArgument(i)->getName() +
353 0 : " is nearly linearly dependent (" + std::to_string(b.gramDropped()) +
354 : " directions left out): KERNEL_DX is too large for SKETCH_NBASIS");
355 : }
356 : }
357 : }
358 16 : if(this->walkers_mpi_) {
359 4 : this->mpi_size_ = getNumberOfReplicas();
360 4 : this->mpi_rank_ = getReplicaIndex();
361 : }
362 :
363 16 : parse("SKETCH_UNTIL", this->sketch_until_);
364 :
365 16 : parseFlag("NONINTRUSIVE", this->nonintrusive_);
366 32 : parseFlag("DETERMINISTIC", this->deterministic_);
367 :
368 16 : if(getRestart()) {
369 : // the last TT in the file is the bias; the hills file holds the hills deposited after it
370 3 : const unsigned ntt = ttCount(this->ttfilename_);
371 3 : if(ntt > 0) {
372 3 : this->sketch_count_ = ntt + 1;
373 3 : setVb(ttRead(this->ttfilename_, this->sketch_count_));
374 3 : if(length(this->vb_) != static_cast<int>(this->d_)) {
375 0 : error("the TT in " + this->ttfilename_ + " has " + std::to_string(length(this->vb_)) +
376 0 : " dimensions, but TTMETAD has " + std::to_string(this->d_) + " arguments");
377 : }
378 9 : for(unsigned i = 0; i < this->d_; ++i) {
379 6 : if(this->vbDense_.nbasis(i) != this->sketch_basis_[i].nbasis()) {
380 0 : error("the TT in " + this->ttfilename_ + " has " + std::to_string(this->vbDense_.nbasis(i)) +
381 0 : " basis functions along argument " + std::to_string(i + 1) + ", but SKETCH_NBASIS gives " +
382 0 : std::to_string(this->sketch_basis_[i].nbasis()));
383 : }
384 : }
385 3 : if(this->nonintrusive_ && !ttMomentsRead(this->ttfilename_, this->sketch_count_, this->d_ - 1, this->B_prev_, this->A_prev_)) {
386 0 : error("NONINTRUSIVE restart needs the tensor moments of the last sketch, which are not in " + this->ttfilename_);
387 : }
388 : }
389 3 : log << " restarting from TT " << this->sketch_count_ - 1 << " in " << this->ttfilename_ << "\n";
390 3 : IFile hills_ifile;
391 3 : hills_ifile.link(*this);
392 3 : if(this->walkers_mpi_) {
393 0 : hills_ifile.enforceSuffix("");
394 : }
395 3 : if(!hills_ifile.FileExist(this->hillsfname_)) {
396 0 : error("The hills file cannot be found");
397 : }
398 3 : hills_ifile.open(this->hillsfname_);
399 3 : readGaussians(&hills_ifile);
400 3 : hills_ifile.close();
401 3 : } else if(writerReplica()) {
402 : // a fresh run: keep a TT file left by a previous run as a backup, as for the other output files,
403 : // so that it cannot be mistaken for the TTs of this run at a later restart
404 11 : OFile backup;
405 11 : backup.link(*this);
406 11 : backup.backupFile("bck", this->ttfilename_);
407 11 : }
408 :
409 16 : log << " TTs are stored in " << this->ttfilename_ << "\n";
410 16 : log << " Bibliography " << plumed.cite("Strand, Yang, Khoo, and Dinner, https://arxiv.org/abs/2603.13549 (2026)");
411 16 : log << plumed.cite("Fishman, White, and Stoudenmire, https://arxiv.org/abs/2007.14822 (2020)");
412 16 : log << "\n";
413 :
414 : if(writerReplica()) {
415 14 : this->hillsOfile_.link(*this);
416 14 : if(this->walkers_mpi_) {
417 : // a single hills file shared by all the walkers (written by the first one);
418 : // otherwise each replica has its own, with the replica suffix
419 4 : this->hillsOfile_.enforceSuffix("");
420 : }
421 14 : this->hillsOfile_.open(this->hillsfname_);
422 14 : if(this->fmt_.length() > 0) {
423 14 : this->hillsOfile_.fmtField(this->fmt_);
424 : }
425 : hillsOfile_.setHeavyFlush();
426 46 : for(unsigned i = 0; i < this->d_; ++i) {
427 32 : hillsOfile_.setupPrintValue(getPntrToArgument(i));
428 : }
429 : }
430 16 : }
431 :
432 3 : void TTMetaD::readGaussians(IFile *ifile) {
433 3 : std::vector<double> center(this->d_);
434 3 : std::vector<double> sigma(this->d_);
435 : double height;
436 3 : int nhills = 0;
437 :
438 3 : std::vector<Value> tmpvalues;
439 9 : for(unsigned j = 0; j < this->d_; ++j) {
440 18 : tmpvalues.push_back(Value(this, getPntrToArgument(j)->getName(), false));
441 : }
442 :
443 5 : while(scanOneHill(ifile, tmpvalues, center, sigma, height)) {
444 2 : ++nhills;
445 2 : if(this->welltemp_ && this->biasf_ > 1.0) {
446 1 : height *= (this->biasf_ - 1.0) / this->biasf_;
447 : }
448 4 : this->hills_.push_back(Gaussian(height, center, sigma));
449 : }
450 3 : log << " " << nhills << " hills retrieved\n";
451 6 : }
452 :
453 50 : void TTMetaD::setVb(const MPS& vb) {
454 50 : this->vb_ = vb;
455 50 : this->vbDense_ = length(vb) == 0 ? DenseTT() : DenseTT(vb);
456 50 : }
457 :
458 376 : void TTMetaD::clampToInterval(std::vector<double>& cv, std::vector<char>& out) const {
459 376 : out.assign(this->d_, 0);
460 1232 : for(unsigned i = 0; i < this->d_; ++i) {
461 856 : if(cv[i] < this->imin_[i]) {
462 0 : cv[i] = this->imin_[i];
463 0 : out[i] = 1;
464 856 : } else if(cv[i] > this->imax_[i]) {
465 8 : cv[i] = this->imax_[i];
466 8 : out[i] = 1;
467 : }
468 : }
469 376 : }
470 :
471 5 : bool TTMetaD::scanOneHill(IFile* ifile, std::vector<Value>& tmpvalues, std::vector<double>& center, std::vector<double>& sigma, double& height) {
472 : double dummy;
473 10 : if(ifile->scanField("time", dummy)) {
474 2 : unsigned ncv = tmpvalues.size();
475 6 : for(unsigned i = 0; i < ncv; ++i) {
476 4 : ifile->scanField(&tmpvalues[i]);
477 4 : if(tmpvalues[i].isPeriodic() && !getPntrToArgument(i)->isPeriodic()) {
478 0 : error("in hills file periodicity for variable " + tmpvalues[i].getName() + " does not match periodicity in input");
479 4 : } else if(tmpvalues[i].isPeriodic()) {
480 : std::string imin, imax;
481 3 : tmpvalues[i].getDomain(imin, imax);
482 : std::string rmin, rmax;
483 3 : getPntrToArgument(i)->getDomain(rmin, rmax);
484 3 : if(imin != rmin || imax != rmax) {
485 0 : error("in hills file periodicity for variable " + tmpvalues[i].getName() + " does not match periodicity in input");
486 : }
487 : }
488 4 : center[i] = tmpvalues[i].get();
489 : }
490 6 : for(unsigned i = 0; i < ncv; ++i) {
491 8 : ifile->scanField("sigma_" + getPntrToArgument(i)->getName(), sigma[i]);
492 : }
493 :
494 2 : ifile->scanField("height", height);
495 2 : ifile->scanField("biasf", dummy);
496 2 : ifile->scanField();
497 2 : return true;
498 : } else {
499 : return false;
500 : }
501 : }
502 :
503 273 : void TTMetaD::writeGaussian(const Gaussian& hill, OFile&file) {
504 273 : file.printField("time", getTimeStep() * getStep());
505 889 : for(unsigned i = 0; i < this->d_; ++i) {
506 616 : file.printField(getPntrToArgument(i), hill.center[i]);
507 : }
508 889 : for(unsigned i = 0; i < this->d_; ++i) {
509 1232 : file.printField("sigma_" + getPntrToArgument(i)->getName(), hill.sigma[i]);
510 : }
511 273 : double height = hill.height;
512 273 : if(this->welltemp_ && this->biasf_ > 1.0) {
513 40 : height *= this->biasf_ / (this->biasf_ - 1.0);
514 : }
515 546 : file.printField("height", height).printField("biasf", this->biasf_);
516 273 : file.printField();
517 273 : }
518 :
519 : // Called every MD step: evaluate total bias (TT + remaining Gaussians) and set forces.
520 336 : void TTMetaD::calculate() {
521 336 : std::vector<double> cv(this->d_);
522 1092 : for(unsigned i = 0; i < this->d_; ++i) {
523 756 : cv[i] = getArgument(i);
524 : }
525 :
526 : // outside the interval the bias is the one at the closest boundary and the force along that CV is zero
527 336 : std::vector<char> out;
528 336 : clampToInterval(cv, out);
529 :
530 336 : std::vector<double> der(this->d_, 0.0);
531 :
532 336 : double ene = getBiasAndDerivatives(cv, der);
533 1092 : for(unsigned i = 0; i < this->d_; ++i) {
534 756 : if(out[i]) {
535 4 : der[i] = 0.0;
536 : }
537 : }
538 336 : setBias(ene);
539 1092 : for(unsigned i = 0; i < this->d_; ++i) {
540 756 : setOutputForce(i, -der[i]);
541 : }
542 672 : }
543 :
544 : // Called every MD step after forces are applied.
545 : // Two separate triggers:
546 : // 1. Every sketch_stride_ steps: run paraSketch() to compress accumulated hills into vb_,
547 : // clear the hills list, and rewind the HILLS file for the next interval.
548 : // 2. Every stride_ steps: deposit a new Gaussian hill and write it to the HILLS file.
549 336 : void TTMetaD::update() {
550 : // SKETCH_UNTIL: at the first step after this time the pending hills are sketched
551 : // one last time and the bias is frozen
552 336 : const bool freezeNow = !this->frozen_ && !this->isFirstStep_ && getTime() >= this->sketch_until_;
553 : // all the walkers hold the same list of hills, so they all take the same decision here
554 336 : const bool nowAddATT = !this->frozen_ && !this->isFirstStep_ && !this->hills_.empty() &&
555 261 : (getStep() % this->sketch_stride_ == 0 || freezeNow);
556 :
557 : if(nowAddATT) {
558 : if(writerReplica()) {
559 43 : this->hillsOfile_.flush();
560 43 : unsigned N = this->hills_.size();
561 43 : log << "Sample limits\n";
562 150 : for(unsigned i = 0; i < this->d_; ++i) {
563 107 : auto [large, small] = this->sketch_basis_[i].dom();
564 107 : unsigned nout = 0;
565 695 : for(unsigned j = 0; j < N; ++j) {
566 588 : const double c = this->hills_[j].center[i];
567 588 : if(c < this->imin_[i] || c > this->imax_[i]) {
568 4 : ++nout;
569 : }
570 588 : large = std::max(large, c);
571 588 : small = std::min(small, c);
572 : }
573 107 : log << small << " " << large << "\n";
574 107 : if(nout > 0) {
575 3 : log << " WARNING: " << nout << " hills are centered outside the interval of argument " << getPntrToArgument(i)->getName()
576 3 : << ": the periodic basis folds them back into the interval\n";
577 : }
578 : }
579 :
580 43 : log << "\nEmpirical means:\n";
581 43 : std::vector<double> muhat(this->d_, 0.0);
582 150 : for(unsigned k = 0; k < this->d_; ++k) {
583 695 : for(unsigned j = 0; j < N; ++j) {
584 588 : muhat[k] += this->hills_[j].center[k] / N;
585 : }
586 107 : log << muhat[k] << " ";
587 : }
588 43 : if(N > 1) {
589 33 : log << "\nEmpirical covariance matrix:\n";
590 33 : Matrix<double> sigmahat(this->d_, this->d_);
591 111 : for(unsigned k = 0; k < this->d_; ++k) {
592 213 : for(unsigned l = k; l < this->d_; ++l) {
593 : double sum = 0.0;
594 1059 : for(unsigned j = 0; j < N; ++j) {
595 924 : sum += (this->hills_[j].center[k] - muhat[k]) * (this->hills_[j].center[l] - muhat[l]);
596 : }
597 135 : sigmahat(k, l) = sigmahat(l, k) = sum / (N - 1);
598 : }
599 : }
600 33 : matrixOut(log, sigmahat);
601 : } else {
602 10 : log << "\n";
603 : }
604 :
605 : // record current bias at hill centers before update, for computing relative error after
606 43 : std::vector<std::vector<double>> x(N);
607 43 : std::vector<double> A0(N);
608 304 : for(unsigned i = 0; i < N; ++i) {
609 261 : x[i] = this->hills_[i].center;
610 261 : A0[i] = getBias(x[i]);
611 : }
612 :
613 43 : log << "\nStarting TT-sketch...\n";
614 43 : log.flush();
615 43 : paraSketch();
616 43 : ++this->sketch_count_;
617 :
618 43 : this->hills_.clear();
619 :
620 : // compute relative L2 approximation error at hill centers: ||V_new - V_old|| / ||V_old||
621 43 : std::vector<double> diff(N);
622 304 : for(unsigned i = 0; i < N; ++i) {
623 261 : diff[i] = getBias(x[i]) - A0[i];
624 : }
625 43 : log << "Relative l2 error = " << sqrt(norm(diff) / norm(A0)) << "\n\n";
626 43 : log.flush();
627 :
628 : // all ranks of the replica hold the same vb_; one of them writes it
629 43 : if(comm.Get_rank() == 0) {
630 41 : ttWrite(this->ttfilename_, this->vb_, this->sketch_count_);
631 41 : if(this->nonintrusive_) {
632 12 : ttMomentsWrite(this->ttfilename_, this->B_prev_, this->A_prev_, this->sketch_count_);
633 : }
634 : }
635 43 : }
636 :
637 47 : if(this->walkers_mpi_) {
638 8 : if(comm.Get_rank() == 0) {
639 4 : multi_sim_comm.Bcast(this->sketch_count_, 0);
640 : }
641 8 : comm.Bcast(this->sketch_count_, 0);
642 8 : if(this->mpi_rank_ != 0) {
643 4 : this->hills_.clear();
644 4 : setVb(ttRead(this->ttfilename_, this->sketch_count_));
645 : }
646 : }
647 : // the hills are now in the TT: the hills file restarts empty
648 : if(writerReplica()) {
649 43 : this->hillsOfile_.rewind();
650 43 : this->hillsOfile_.clearFields();
651 43 : if(this->fmt_.length() > 0) {
652 43 : this->hillsOfile_.fmtField(this->fmt_);
653 : }
654 : hillsOfile_.setHeavyFlush();
655 150 : for(unsigned i = 0; i < this->d_; ++i) {
656 107 : hillsOfile_.setupPrintValue(getPntrToArgument(i));
657 : }
658 : }
659 : }
660 336 : if(freezeNow) {
661 2 : this->frozen_ = true;
662 2 : log << "TT bias frozen at time " << getTime() << "\n";
663 : }
664 :
665 : bool nowAddAHill;
666 336 : if(getStep() % this->stride_ == 0 && !isFirstStep_ && !this->frozen_) {
667 : nowAddAHill = true;
668 : } else {
669 : nowAddAHill = false;
670 63 : this->isFirstStep_ = false;
671 : }
672 :
673 336 : std::vector<double> cv(this->d_);
674 1092 : for(unsigned i = 0; i < this->d_; ++i) {
675 756 : cv[i] = getArgument(i);
676 : }
677 :
678 336 : if(nowAddAHill) {
679 273 : double height = getHeight(cv);
680 :
681 273 : if(this->walkers_mpi_) {
682 80 : std::vector<double> all_cv(this->mpi_size_ * this->d_, 0.0);
683 80 : std::vector<double> all_sigma(this->mpi_size_ * this->sigma0_.size(), 0.0);
684 80 : std::vector<double> all_height(this->mpi_size_, 0.0);
685 80 : if(comm.Get_rank() == 0) {
686 40 : multi_sim_comm.Allgather(cv, all_cv);
687 40 : multi_sim_comm.Allgather(this->sigma0_, all_sigma);
688 40 : multi_sim_comm.Allgather(height * (this->biasf_ > 1.0 ? this->biasf_ / (this->biasf_ - 1.0) : 1.0), all_height);
689 : }
690 80 : comm.Bcast(all_cv, 0);
691 80 : comm.Bcast(all_sigma, 0);
692 80 : comm.Bcast(all_height, 0);
693 :
694 240 : for(int i = 0; i < this->mpi_size_; i++) {
695 160 : std::vector<double> cv_now(this->d_);
696 160 : std::vector<double> sigma_now(this->sigma0_.size());
697 480 : for(unsigned j = 0; j < this->d_; j++) {
698 320 : cv_now[j] = all_cv[i * this->d_ + j];
699 : }
700 480 : for(unsigned j = 0; j < this->sigma0_.size(); j++) {
701 320 : sigma_now[j] = all_sigma[i * this->sigma0_.size() + j];
702 : }
703 160 : double fact = (this->biasf_ > 1.0 ? (this->biasf_ - 1.0) / this->biasf_ : 1.0);
704 160 : Gaussian newhill(all_height[i] * fact, cv_now, sigma_now);
705 160 : this->hills_.push_back(newhill);
706 160 : if(this->mpi_rank_ == 0) {
707 80 : writeGaussian(newhill, hillsOfile_);
708 : }
709 160 : }
710 : } else {
711 193 : Gaussian newhill(height, cv, this->sigma0_);
712 193 : this->hills_.push_back(newhill);
713 193 : writeGaussian(newhill, hillsOfile_);
714 193 : }
715 : }
716 :
717 336 : if(getStep() % this->sketch_stride_ == 1 && !this->frozen_) {
718 0 : log << "Vbias update " << this->sketch_count_ << "...\n\n";
719 0 : log.flush();
720 : }
721 336 : }
722 :
723 273 : double TTMetaD::getHeight(const std::vector<double>& cv) {
724 273 : double height = this->height0_;
725 273 : if(this->welltemp_) {
726 : // the bias felt by the system (see calculate())
727 40 : std::vector<double> cvc(cv);
728 40 : std::vector<char> out;
729 40 : clampToInterval(cvc, out);
730 40 : double vbias = getBias(cvc);
731 40 : if(this->biasf_ > 1.0) {
732 40 : height = this->height0_ * exp(-vbias / (this->kbt_ * (this->biasf_ - 1.0)));
733 : } else {
734 0 : height = this->height0_ * exp(-vbias / this->kbt_);
735 : }
736 40 : }
737 273 : return height;
738 : }
739 :
740 : // Total bias = TT approximation of accumulated hills + sum of recently deposited Gaussians.
741 562 : double TTMetaD::getBias(const std::vector<double>& cv) {
742 562 : double bias = this->vbDense_.empty() ? 0.0 : this->vbDense_.eval(this->sketch_basis_, cv, this->sketch_conv_, nullptr);
743 562 : unsigned nt = OpenMP::getNumThreads();
744 562 : #pragma omp parallel num_threads(nt)
745 : {
746 : #pragma omp for reduction(+:bias) nowait
747 : for(unsigned i = 0; i < hills_.size(); ++i) {
748 : bias += evaluateGaussian(cv, this->hills_[i]);
749 : }
750 : }
751 562 : return bias;
752 : }
753 :
754 336 : double TTMetaD::getBiasAndDerivatives(const std::vector<double>& cv, std::vector<double>& der) {
755 : double bias = 0.0;
756 336 : if (!this->vbDense_.empty()) {
757 227 : bias = this->vbDense_.eval(this->sketch_basis_, cv, this->sketch_conv_, &der);
758 : } else {
759 109 : der.assign(this->d_, 0.0);
760 : }
761 336 : unsigned nt = OpenMP::getNumThreads();
762 336 : if(this->hills_.size() < 2 * nt || nt == 1) {
763 179 : std::vector<double> dp(this->d_);
764 378 : for(unsigned i = 0; i < this->hills_.size(); ++i) {
765 199 : bias += evaluateGaussianAndDerivatives(cv, this->hills_[i], der, dp);
766 : }
767 : } else {
768 157 : #pragma omp parallel num_threads(nt)
769 : {
770 : std::vector<double> omp_deriv(this->d_, 0.0);
771 : std::vector<double> dp(this->d_);
772 : #pragma omp for reduction(+:bias) nowait
773 : for(unsigned i = 0; i < this->hills_.size(); ++i) {
774 : bias += evaluateGaussianAndDerivatives(cv, this->hills_[i], omp_deriv, dp);
775 : }
776 : #pragma omp critical
777 : for(unsigned i = 0; i < this->d_; ++i) {
778 : der[i] += omp_deriv[i];
779 : }
780 : }
781 : }
782 336 : return bias;
783 : }
784 :
785 2795 : double TTMetaD::evaluateGaussian(const std::vector<double>& cv, const Gaussian& hill) {
786 : double dp2 = 0.0;
787 8703 : for(unsigned i = 0; i < this->d_; i++) {
788 5908 : double dp = difference(i, hill.center[i], cv[i]) * hill.invsigma[i];
789 5908 : dp2 += dp * dp;
790 : }
791 2795 : dp2 *= 0.5;
792 :
793 : double bias = 0.0;
794 2795 : if(dp2 < dp2cutoff) {
795 2049 : bias = hill.height * exp(-dp2);
796 : }
797 :
798 2795 : return bias;
799 : }
800 :
801 1505 : double TTMetaD::evaluateGaussianAndDerivatives(const std::vector<double>& cv, const Gaussian& hill, std::vector<double>& der, std::vector<double>& dp) {
802 : double dp2 = 0.0;
803 : double bias = 0.0;
804 4689 : for(unsigned i = 0; i < this->d_; i++) {
805 3184 : dp[i] = difference(i, hill.center[i], cv[i]) * hill.invsigma[i];
806 3184 : dp2 += dp[i] * dp[i];
807 : }
808 1505 : dp2 *= 0.5;
809 1505 : if(dp2 < dp2cutoff) {
810 1090 : bias = hill.height * exp(-dp2);
811 3369 : for(unsigned i = 0; i < this->d_; i++) {
812 2279 : der[i] -= bias * dp[i] * hill.invsigma[i];
813 : }
814 : }
815 :
816 1505 : return bias;
817 : }
818 :
819 : // TT-Sketch algorithm: compresses the N accumulated Gaussian hills plus the previous
820 : // TT bias vb_ into a new low-rank TT approximation G, which becomes the updated vb_.
821 : //
822 : // Overview (see paper Algorithm 2):
823 : // 1. createTTCoeff(): build a random sketch coefficient TT `coeff` of rank sketch_rc_.
824 : // 2. intBasisSample(): compute M[i](j,k) = <phi_k, g_j>_{dim i}, the inner product of
825 : // each basis function with each Gaussian along dimension i.
826 : // 3. formTensorMoment(): contract coeff with M to get the tensor moment Bemp and the
827 : // cross-gram matrices A of the bonds (normal equations).
828 : // 4. For each bond between cores k and k+1:
829 : // - Add the contribution of the previous bias (or of the previous moments) to A.
830 : // - SVD A to get the rank-trimmed projection V (truncating to sketch_r_ or sketch_cutoff_).
831 : // 5. Recover TT cores G from Bemp and the pseudo-inverse of U*S.
832 : // 6. Apply Gram matrix correction (kernel basis only) to convert to dual basis coefficients.
833 43 : void TTMetaD::paraSketch() {
834 43 : auto coeff = createTTCoeff();
835 43 : auto M = intBasisSample();
836 43 : MPS G(this->d_);
837 :
838 43 : auto [Bemp, AMats] = formTensorMoment(M, coeff);
839 43 : MPS Bemp_Vb;
840 43 : std::vector<ITensor> envi_L_Vb;
841 43 : std::vector<ITensor> envi_R_Vb;
842 43 : if(this->sketch_count_ != 1) {
843 32 : if(this->nonintrusive_) {
844 : // Non-intrusive variant: accumulate tensor moments directly across sketches.
845 : // B_prev_ stores the sum of Bemp from all prior intervals; re-index its bond/site
846 : // indices to match the current coeff TT's indices before adding.
847 10 : auto sites = siteInds(Bemp);
848 10 : auto sites_prev = siteInds(this->B_prev_);
849 10 : auto links = linkInds(Bemp);
850 10 : auto links_prev = linkInds(this->B_prev_);
851 33 : for(unsigned i = 1; i <= this->d_; ++i) {
852 23 : this->B_prev_.ref(i) *= delta(sites(i), sites_prev(i));
853 23 : if(i != 1) {
854 13 : this->B_prev_.ref(i) *= delta(links(i - 1), links_prev(i - 1));
855 : }
856 23 : if (i != this->d_) {
857 13 : this->B_prev_.ref(i) *= delta(links(i), links_prev(i));
858 : }
859 23 : Bemp.ref(i) += this->B_prev_(i);
860 : }
861 : } else {
862 : // Intrusive variant: project the previous TT bias vb_ onto the current sketch's
863 : // basis by pre-multiplying each core by the Gram matrix (kernel basis) or the
864 : // convolution inner-product matrix (Fourier conv mode), then add its tensor moment.
865 22 : if(this->sketch_basis_[0].kernel()) {
866 : // Convert vb_ from dual basis to standard basis by applying Gram matrix per core
867 32 : for(unsigned i = 1; i <= this->d_; ++i) {
868 24 : auto s = siteIndex(this->vb_, i);
869 24 : ITensor gram(s, prime(s));
870 216 : for(int j = 1; j <= dim(s); ++j) {
871 1728 : for(int l = 1; l <= dim(s); ++l) {
872 4608 : gram.set(s = j, prime(s) = l, this->sketch_basis_[i - 1].gram()(j - 1, l - 1));
873 : }
874 : }
875 24 : this->vb_.ref(i) *= gram;
876 24 : this->vb_.ref(i).noPrime();
877 24 : }
878 : }
879 : // Note: in SKETCH_WIDTH (conv) mode the coefficients of vb_ are those of the unsmoothed bias,
880 : // the smoothing being applied when the basis is evaluated, so they are added as they are.
881 : // Damping them here as well (as done before) smoothed the old bias again at every sketch,
882 : // so that the intrusive variant did not match NONINTRUSIVE even when the TT is exact.
883 22 : auto vbresult = formTensorMomentVb(coeff);
884 22 : Bemp_Vb = std::get<0>(vbresult);
885 22 : envi_L_Vb = std::get<1>(vbresult);
886 22 : envi_R_Vb = std::get<2>(vbresult);
887 80 : for(unsigned i = 1; i <= this->d_; ++i) {
888 58 : Bemp.ref(i) += Bemp_Vb(i);
889 : }
890 : }
891 : }
892 43 : if(this->nonintrusive_) {
893 12 : this->B_prev_ = Bemp;
894 12 : if(this->sketch_count_ == 1) {
895 2 : this->A_prev_.resize(this->d_ - 1);
896 : }
897 : }
898 : // For each bond between cores k-1 and k, form the cross-gram matrix A = L^T * R,
899 : // where L = envi_L[k] and R = envi_R[k-2] are the left/right environment projections
900 : // of coeff onto the sample indices. SVD A to find the rank-trimmed subspace V.
901 43 : auto links = linkInds(coeff);
902 43 : std::vector<ITensor> U(this->d_), S(this->d_), V(this->d_);
903 43 : std::vector<Index> links_trimmed;
904 107 : for(unsigned core_id = 2; core_id <= this->d_; ++core_id) {
905 64 : int rank = dim(links(core_id - 1));
906 64 : const Matrix<double>& AMat = AMats[core_id - 2];
907 64 : Matrix<double> LMat, RMat, Lt, AMat_Vb;
908 :
909 64 : if(this->sketch_count_ != 1 && !this->nonintrusive_) {
910 36 : auto ivb = linkIndex(this->vb_, core_id - 1);
911 36 : int rank_vb = dim(ivb);
912 36 : LMat = Matrix<double>(rank_vb, rank);
913 36 : RMat = Matrix<double>(rank_vb, rank);
914 406 : for(int i = 1; i <= rank_vb; ++i) {
915 18438 : for(int j = 1; j <= rank; ++j) {
916 18068 : LMat(i - 1, j - 1) = envi_L_Vb[core_id - 1].elt(ivb = i, links(core_id - 1) = j);
917 18068 : RMat(i - 1, j - 1) = envi_R_Vb[core_id - 2].elt(ivb = i, links(core_id - 1) = j);
918 : }
919 : }
920 36 : transpose(LMat, Lt);
921 36 : mult(Lt, RMat, AMat_Vb);
922 : }
923 :
924 64 : ITensor A(prime(links(core_id - 1)), links(core_id - 1));
925 1288 : for(int i = 1; i <= rank; ++i) {
926 60312 : for(int j = 1; j <= rank; ++j) {
927 118176 : A.set(prime(links(core_id - 1)) = i, links(core_id - 1) = j, AMat(i - 1, j - 1));
928 : }
929 : }
930 64 : if(this->sketch_count_ != 1) {
931 49 : if(this->nonintrusive_) {
932 13 : this->A_prev_[core_id - 2] *= delta(this->A_prev_[core_id - 2].index(1), A.index(1));
933 13 : this->A_prev_[core_id - 2] *= delta(this->A_prev_[core_id - 2].index(2), A.index(2));
934 13 : A += this->A_prev_[core_id - 2];
935 : } else {
936 36 : ITensor A_Vb(prime(links(core_id - 1)), links(core_id - 1));
937 562 : for(int i = 1; i <= rank; ++i) {
938 23066 : for(int j = 1; j <= rank; ++j) {
939 45080 : A_Vb.set(prime(links(core_id - 1)) = i, links(core_id - 1) = j, AMat_Vb(i - 1, j - 1));
940 : }
941 : }
942 36 : A += A_Vb;
943 36 : }
944 : }
945 64 : if(this->nonintrusive_) {
946 16 : this->A_prev_[core_id - 2] = A;
947 : }
948 64 : auto original_link_tags = tags(links(core_id - 1));
949 64 : V[core_id - 1] = ITensor(links(core_id - 1));
950 64 : if(this->sketch_r_ > 0) {
951 112 : svd(A, U[core_id - 1], S[core_id - 1], V[core_id - 1],
952 56 : {"Cutoff=", this->sketch_cutoff_, "RightTags=", original_link_tags, "MaxDim=", this->sketch_r_});
953 : } else {
954 16 : svd(A, U[core_id - 1], S[core_id - 1], V[core_id - 1], {"Cutoff=", this->sketch_cutoff_, "RightTags=", original_link_tags});
955 : }
956 64 : links_trimmed.push_back(commonIndex(S[core_id - 1], V[core_id - 1]));
957 64 : }
958 :
959 : // Recover TT cores G from Bemp and the trimmed subspaces.
960 : // G[1] = Bemp[1] * V[1] (project first core onto rank-trimmed subspace)
961 : // G[k] = pinv(U[k]*S[k]) * Bemp[k] * V[k] for 2 <= k < d
962 : // G[d] = pinv(U[d]*S[d]) * Bemp[d]
963 : // Since U has orthonormal columns, pinv(U*S) = S^+ * U^T, where S^+ inverts only the
964 : // singular values above a relative threshold: the others (e.g. all but one when a single
965 : // hill is sketched, so that A has rank one) are rounding errors, and dividing Bemp by them
966 : // would fill the TT with noise of order one.
967 43 : G.ref(1) = Bemp(1) * V[1];
968 43 : std::vector<int> effective_rank(this->d_, 0);
969 107 : for(unsigned core_id = 2; core_id <= this->d_; ++core_id) {
970 64 : int rank = dim(links(core_id - 1)), rank_trimmed = dim(links_trimmed[core_id - 2]);
971 64 : const Index iu = commonIndex(U[core_id - 1], S[core_id - 1]);
972 64 : const Index iv = links_trimmed[core_id - 2];
973 64 : std::vector<double> sinv(rank_trimmed, 0.0);
974 64 : double smax = 0.0;
975 971 : for(int j = 1; j <= rank_trimmed; ++j) {
976 971 : smax = std::max(smax, S[core_id - 1].elt(iu = j, iv = j));
977 : }
978 971 : for(int j = 1; j <= rank_trimmed; ++j) {
979 1814 : const double sj = S[core_id - 1].elt(iu = j, iv = j);
980 907 : if(sj > svdTolerance * smax) {
981 216 : sinv[j - 1] = 1.0 / sj;
982 216 : ++effective_rank[core_id - 2];
983 : }
984 : }
985 64 : ITensor Pinv(iv, links(core_id - 1));
986 971 : for(int i = 1; i <= rank_trimmed; ++i) {
987 48637 : for(int j = 1; j <= rank; ++j) {
988 143190 : Pinv.set(iv = i, links(core_id - 1) = j, sinv[i - 1] * U[core_id - 1].elt(prime(links(core_id - 1)) = j, iu = i));
989 : }
990 : }
991 64 : G.ref(core_id) = Pinv * Bemp(core_id);
992 64 : if(core_id != this->d_) {
993 21 : G.ref(core_id) *= V[core_id];
994 : }
995 64 : }
996 :
997 43 : log << "Final ranks ";
998 107 : for(unsigned i = 1; i < this->d_; ++i) {
999 128 : log << dim(linkIndex(G, i)) << " ";
1000 : }
1001 43 : log << "(effective ";
1002 107 : for(unsigned i = 1; i < this->d_; ++i) {
1003 64 : log << effective_rank[i - 1] << " ";
1004 : }
1005 43 : log << ")\n";
1006 43 : log.flush();
1007 :
1008 43 : if(this->sketch_basis_[0].kernel()) {
1009 : // Apply pseudo-inverse of the Gram matrix to each core to convert from the
1010 : // primal Gram representation back to dual basis coefficients (G^+ * G_core).
1011 36 : for(unsigned i = 1; i <= this->d_; ++i) {
1012 27 : auto s = siteIndex(G, i);
1013 27 : ITensor ginv(s, prime(s));
1014 243 : for(int j = 1; j <= dim(s); ++j) {
1015 1944 : for(int l = 1; l <= dim(s); ++l) {
1016 5184 : ginv.set(s = j, prime(s) = l, this->sketch_basis_[i - 1].ginv()(j - 1, l - 1));
1017 : }
1018 : }
1019 27 : G.ref(i) *= ginv;
1020 27 : G.ref(i).noPrime();
1021 27 : }
1022 : }
1023 :
1024 43 : setVb(G);
1025 129 : }
1026 :
1027 : // Build the random sketch coefficient TT `coeff` of bond dimension sketch_rc_.
1028 : // Each element is drawn from N(0,1); then each core is multiplied by a diagonal
1029 : // matrix diag(1, alpha, alpha, ...) so the constant basis function (pos=1) has
1030 : // full weight while all other harmonics are scaled down by sketch_alpha_. This
1031 : // ensures the sketch captures the large constant component of the bias accurately.
1032 43 : MPS TTMetaD::createTTCoeff() const {
1033 43 : unsigned seed = this->deterministic_ ? 42u : static_cast<unsigned int>(time(nullptr));
1034 : // all ranks of a replica must build the same sketch
1035 43 : comm.Bcast(seed, 0);
1036 43 : std::default_random_engine generator(seed);
1037 : std::normal_distribution<double> distribution(0.0, 1.0);
1038 : int n = this->sketch_basis_[0].nbasis();
1039 43 : auto sites = SiteSet(this->d_, n);
1040 43 : auto coeff = MPS(sites, this->sketch_rc_);
1041 329 : for(int j = 1; j <= n; ++j) {
1042 4906 : for(int k = 1; k <= this->sketch_rc_; ++k) {
1043 18480 : coeff.ref(1).set(sites(1) = j, linkIndex(coeff, 1) = k, distribution(generator));
1044 : }
1045 : }
1046 64 : for(unsigned i = 2; i <= this->d_ - 1; ++i) {
1047 177 : for(int j = 1; j <= n; ++j) {
1048 4116 : for(int k = 1; k <= this->sketch_rc_; ++k) {
1049 209160 : for(int l = 1; l <= this->sketch_rc_; ++l) {
1050 1026000 : coeff.ref(i).set(sites(i) = j, linkIndex(coeff, i - 1) = k, linkIndex(coeff, i) = l, distribution(generator));
1051 : }
1052 : }
1053 : }
1054 : }
1055 329 : for(int j = 1; j <= n; ++j) {
1056 4906 : for(int k = 1; k <= this->sketch_rc_; ++k) {
1057 18480 : coeff.ref(this->d_).set(sites(this->d_) = j, linkIndex(coeff, this->d_ - 1) = k, distribution(generator));
1058 : }
1059 : }
1060 150 : for(unsigned i = 1; i <= this->d_; ++i) {
1061 107 : auto s = sites(i);
1062 107 : auto sp = prime(s);
1063 107 : std::vector<double> Avec(n, this->sketch_alpha_);
1064 107 : Avec[0] = 1.0;
1065 107 : auto A = diagITensor(Avec, s, sp);
1066 107 : coeff.ref(i) *= A;
1067 107 : coeff.ref(i).noPrime();
1068 107 : }
1069 43 : return coeff;
1070 0 : }
1071 :
1072 : // Compute the matrix M[dim i][sample j, basis k] = <phi_k, g_j>_{dim i}, the
1073 : // analytical inner product of each 1D basis function with the 1D marginal of each
1074 : // Gaussian hill along dimension i. The d-dimensional Gaussian factorizes as a
1075 : // product of d 1D Gaussians, so we distribute height as h^(1/d) per dimension.
1076 : // These integrals are computed in closed form:
1077 : // - Kernel basis: convolution of the kernel function with the 1D Gaussian hill
1078 : // - Fourier basis: Fourier transform of the 1D Gaussian (exponential decay * harmonic)
1079 : // Returns M as a vector of ITensors (one per dimension) and a new IndexSet is_new
1080 : // where is_new(i) indexes the N sample points for dimension i.
1081 43 : std::vector<Matrix<double>> TTMetaD::intBasisSample() const {
1082 43 : unsigned N = this->hills_.size();
1083 : int nb = this->sketch_basis_[0].nbasis();
1084 43 : std::vector<Matrix<double>> M(this->d_, Matrix<double>(N, nb));
1085 150 : for(unsigned i = 1; i <= this->d_; ++i) {
1086 107 : double L = (this->sketch_basis_[i - 1].dom().second - this->sketch_basis_[i - 1].dom().first) / 2;
1087 107 : double a = (this->sketch_basis_[i - 1].dom().second + this->sketch_basis_[i - 1].dom().first) / 2;
1088 695 : for(unsigned j = 1; j <= N; ++j) {
1089 588 : double x = this->hills_[j - 1].center[i - 1];
1090 588 : double w = this->hills_[j - 1].sigma[i - 1];
1091 588 : double h = pow(this->hills_[j - 1].height, 1.0 / this->d_); // per-dimension d-th root of height
1092 4119 : for(int pos = 1; pos <= nb; ++pos) {
1093 : double result = 0.0;
1094 3531 : if(this->sketch_basis_[0].kernel()) {
1095 216 : if(pos == 1) {
1096 : // <1, g_j> = integral of 1D Gaussian = sqrt(2*pi)*sigma
1097 27 : result = h * sqrt(2 * M_PI) * w;
1098 : } else {
1099 : // <kernel_c, g_j> = conv of two Gaussians (widths dx and w), with periodic images
1100 : double c = this->sketch_basis_[i - 1].center(pos - 1);
1101 : double dx = this->sketch_basis_[i - 1].dx();
1102 756 : for(int k = -1; k <= 1; ++k) {
1103 567 : result += exp(-pow(x - c + 2 * k * L, 2) / (2 * (pow(dx, 2) + pow(w, 2)))) * h * sqrt(2 * M_PI) * w /
1104 567 : (sqrt(1 / pow(dx, 2) + 1 / pow(w, 2)) * w);
1105 : }
1106 : }
1107 : } else {
1108 3315 : if(pos == 1) {
1109 : // <1/sqrt(2L), g_j> = sqrt(pi/L) * sigma (Gaussian times constant)
1110 561 : result = h * sqrt(M_PI / L) * w;
1111 2754 : } else if(pos % 2 == 0) {
1112 : // <cos harmonic, g_j>: Fourier transform of Gaussian times exponential damping
1113 1377 : result = exp(-pow(M_PI * w * (pos / 2), 2) / (2 * pow(L, 2))) * h * sqrt(2 * M_PI / L) * w * cos(M_PI * (x - a) * (pos / 2) / L);
1114 : } else {
1115 1377 : result = exp(-pow(M_PI * w * (pos / 2), 2) / (2 * pow(L, 2))) * h * sqrt(2 * M_PI / L) * w * sin(M_PI * (x - a) * (pos / 2) / L);
1116 : }
1117 : }
1118 3531 : M[i - 1](j - 1, pos - 1) = result;
1119 : }
1120 : }
1121 : }
1122 43 : return M;
1123 : }
1124 :
1125 : // Compute the tensor moment B of the new Gaussian hills and the cross-gram matrices A.
1126 : //
1127 : // With K_i(j) = sum_p coeff_i[:, p, :] M_i(j, p) the core i of coeff contracted with the
1128 : // basis integrals of hill j (a r_{i-1} x r_i matrix), the left and right environments are
1129 : // EL_i(j, :) = K_1(j) ... K_i(j) (i = 1..d-1, size r_i)
1130 : // ER_i(j, :) = K_{i+2}(j) ... K_d(j) (i = 0..d-2, size r_{i+1})
1131 : // The tensor moment is
1132 : // B_1[p, b] = sum_j ER_0(j, b) M_1(j, p)
1133 : // B_k[a, p, b] = sum_j EL_{k-1}(j, a) ER_{k-1}(j, b) M_k(j, p) (1 < k < d)
1134 : // B_d[a, p] = sum_j EL_{d-1}(j, a) M_d(j, p)
1135 : // and the cross-gram matrix of bond k = 1..d-1 is A_k = EL_k^T ER_{k-1}.
1136 : // Everything is done on plain arrays (the environments are independent across hills),
1137 : // only B is converted to ITensors at the end.
1138 43 : std::pair<MPS, std::vector<Matrix<double>>> TTMetaD::formTensorMoment(const std::vector<Matrix<double>>& M, const MPS& coeff) const {
1139 43 : const unsigned d = this->d_;
1140 : const unsigned N = M[0].nrows();
1141 43 : const DenseTT C(coeff);
1142 43 : const auto sites = siteInds(coeff);
1143 43 : const auto links = linkInds(coeff);
1144 43 : std::vector<int> r(d + 1), n(d);
1145 193 : for(unsigned i = 0; i <= d; ++i) {
1146 150 : r[i] = C.rank(i);
1147 : }
1148 150 : for(unsigned i = 0; i < d; ++i) {
1149 107 : n[i] = C.nbasis(i);
1150 : }
1151 : // EL[i] for i = 1..d-1 (N x r_i), ER[i] for i = 0..d-2 (N x r_{i+1}), row-major
1152 43 : std::vector<std::vector<double>> EL(d), ER(d - 1);
1153 107 : for(unsigned i = 1; i < d; ++i) {
1154 64 : EL[i].assign(static_cast<size_t>(N) * r[i], 0.0);
1155 64 : ER[i - 1].assign(static_cast<size_t>(N) * r[i], 0.0);
1156 : }
1157 43 : const unsigned nt = OpenMP::getNumThreads();
1158 43 : #pragma omp parallel num_threads(nt)
1159 : {
1160 : std::vector<std::vector<double>> K(d);
1161 : std::vector<double> tmp;
1162 : #pragma omp for schedule(static)
1163 : for(unsigned j = 0; j < N; ++j) {
1164 : // contracted cores K[i] = K_{i+1}(j), r_i x r_{i+1}
1165 : for(unsigned i = 0; i < d; ++i) {
1166 : const int rl = r[i], rr = r[i + 1], nb = n[i];
1167 : K[i].assign(static_cast<size_t>(rl) * rr, 0.0);
1168 : const double* c = C.core(i).data();
1169 : for(int a = 0; a < rl; ++a) {
1170 : double* k = K[i].data() + static_cast<size_t>(a) * rr;
1171 : for(int p = 0; p < nb; ++p) {
1172 : const double m = M[i](j, p);
1173 : const double* cp = c + (static_cast<size_t>(a) * nb + p) * rr;
1174 : for(int b = 0; b < rr; ++b) {
1175 : k[b] += cp[b] * m;
1176 : }
1177 : }
1178 : }
1179 : }
1180 : // left environments
1181 : for(int b = 0; b < r[1]; ++b) {
1182 : EL[1][static_cast<size_t>(j) * r[1] + b] = K[0][b];
1183 : }
1184 : for(unsigned i = 2; i < d; ++i) {
1185 : const int rl = r[i - 1], rr = r[i];
1186 : const double* prev = EL[i - 1].data() + static_cast<size_t>(j) * rl;
1187 : double* out = EL[i].data() + static_cast<size_t>(j) * rr;
1188 : for(int a = 0; a < rl; ++a) {
1189 : const double* k = K[i - 1].data() + static_cast<size_t>(a) * rr;
1190 : for(int b = 0; b < rr; ++b) {
1191 : out[b] += prev[a] * k[b];
1192 : }
1193 : }
1194 : }
1195 : // right environments
1196 : for(int a = 0; a < r[d - 1]; ++a) {
1197 : ER[d - 2][static_cast<size_t>(j) * r[d - 1] + a] = K[d - 1][a];
1198 : }
1199 : for(int i = static_cast<int>(d) - 3; i >= 0; --i) {
1200 : // ER[i] = K_{i+2} ER[i+1], K_{i+2} = K[i+1] is r_{i+1} x r_{i+2}
1201 : const int rl = r[i + 1], rr = r[i + 2];
1202 : const double* next = ER[i + 1].data() + static_cast<size_t>(j) * rr;
1203 : double* out = ER[i].data() + static_cast<size_t>(j) * rl;
1204 : for(int a = 0; a < rl; ++a) {
1205 : const double* k = K[i + 1].data() + static_cast<size_t>(a) * rr;
1206 : double sum = 0.0;
1207 : for(int b = 0; b < rr; ++b) {
1208 : sum += k[b] * next[b];
1209 : }
1210 : out[a] = sum;
1211 : }
1212 : }
1213 : }
1214 : }
1215 :
1216 : // cross-gram matrices
1217 43 : std::vector<Matrix<double>> A(d - 1);
1218 107 : for(unsigned k = 1; k < d; ++k) {
1219 64 : const int rk = r[k];
1220 64 : A[k - 1] = Matrix<double>(rk, rk);
1221 64 : #pragma omp parallel for num_threads(nt) schedule(static)
1222 : for(int a = 0; a < rk; ++a) {
1223 : for(int b = 0; b < rk; ++b) {
1224 : double sum = 0.0;
1225 : for(unsigned j = 0; j < N; ++j) {
1226 : sum += EL[k][static_cast<size_t>(j) * rk + a] * ER[k - 1][static_cast<size_t>(j) * rk + b];
1227 : }
1228 : A[k - 1](a, b) = sum;
1229 : }
1230 : }
1231 : }
1232 :
1233 : // tensor moment
1234 43 : MPS B(d);
1235 : {
1236 43 : const int nb = n[0], rr = r[1];
1237 43 : ITensor core(sites(1), links(1));
1238 329 : for(int p = 0; p < nb; ++p) {
1239 4906 : for(int b = 0; b < rr; ++b) {
1240 4620 : double sum = 0.0;
1241 26634 : for(unsigned j = 0; j < N; ++j) {
1242 22014 : sum += ER[0][static_cast<size_t>(j) * rr + b] * M[0](j, p);
1243 : }
1244 9240 : core.set(sites(1) = p + 1, links(1) = b + 1, sum);
1245 : }
1246 : }
1247 43 : B.ref(1) = core;
1248 43 : }
1249 64 : for(unsigned k = 2; k < d; ++k) {
1250 21 : const int rl = r[k - 1], rr = r[k], nb = n[k - 1];
1251 21 : std::vector<double> Bk(static_cast<size_t>(rl) * rr * nb, 0.0); // [a][b][p]
1252 21 : #pragma omp parallel for num_threads(nt) schedule(static)
1253 : for(int a = 0; a < rl; ++a) {
1254 : for(unsigned j = 0; j < N; ++j) {
1255 : const double la = EL[k - 1][static_cast<size_t>(j) * rl + a];
1256 : const double* er = ER[k - 1].data() + static_cast<size_t>(j) * rr;
1257 : for(int b = 0; b < rr; ++b) {
1258 : const double w = la * er[b];
1259 : double* out = Bk.data() + (static_cast<size_t>(a) * rr + b) * nb;
1260 : for(int p = 0; p < nb; ++p) {
1261 : out[p] += w * M[k - 1](j, p);
1262 : }
1263 : }
1264 : }
1265 : }
1266 21 : ITensor core(links(k - 1), sites(k), links(k));
1267 579 : for(int a = 0; a < rl; ++a) {
1268 29826 : for(int b = 0; b < rr; ++b) {
1269 234468 : for(int p = 0; p < nb; ++p) {
1270 410400 : core.set(links(k - 1) = a + 1, sites(k) = p + 1, links(k) = b + 1, Bk[(static_cast<size_t>(a) * rr + b) * nb + p]);
1271 : }
1272 : }
1273 : }
1274 21 : B.ref(k) = core;
1275 21 : }
1276 : {
1277 43 : const int nb = n[d - 1], rl = r[d - 1];
1278 43 : ITensor core(links(d - 1), sites(d));
1279 709 : for(int a = 0; a < rl; ++a) {
1280 5286 : for(int p = 0; p < nb; ++p) {
1281 4620 : double sum = 0.0;
1282 26634 : for(unsigned j = 0; j < N; ++j) {
1283 22014 : sum += EL[d - 1][static_cast<size_t>(j) * rl + a] * M[d - 1](j, p);
1284 : }
1285 9240 : core.set(links(d - 1) = a + 1, sites(d) = p + 1, sum);
1286 : }
1287 : }
1288 43 : B.ref(d) = core;
1289 43 : }
1290 43 : return std::make_pair(B, A);
1291 86 : }
1292 :
1293 : // Compute the tensor moment B and environments for the previous TT bias vb_.
1294 : // This is the intrusive variant's contribution from the prior accumulated bias.
1295 : // Analogous to formTensorMoment but uses vb_ in place of M*coeff: the "samples"
1296 : // are the vb_ bond indices, and the environments are formed by contracting coeff
1297 : // with vb_ core by core. The result has the same structure as Bemp from the new hills.
1298 22 : std::tuple<MPS, std::vector<ITensor>, std::vector<ITensor>> TTMetaD::formTensorMomentVb(const MPS& coeff) {
1299 : // align vb_ site indices with coeff's site indices for contraction
1300 80 : for(unsigned i = 1; i <= this->d_; ++i) {
1301 116 : this->vb_.ref(i) *= delta(siteIndex(this->vb_, i), siteIndex(coeff, i));
1302 : }
1303 22 : std::vector<ITensor> envi_L(this->d_);
1304 22 : envi_L[1] = coeff(1) * this->vb_(1);
1305 36 : for(unsigned i = 2; i < this->d_; ++i) {
1306 14 : envi_L[i] = envi_L[i - 1] * coeff(i) * this->vb_(i);
1307 : }
1308 :
1309 22 : std::vector<ITensor> envi_R(this->d_);
1310 22 : envi_R[this->d_ - 2] = coeff(this->d_) * this->vb_(this->d_);
1311 36 : for(int i = this->d_ - 3; i >= 0; --i) {
1312 14 : envi_R[i] = envi_R[i + 1] * coeff(i + 2) * this->vb_(i + 2);
1313 : }
1314 :
1315 22 : MPS B(this->d_);
1316 22 : B.ref(1) = this->vb_(1) * envi_R[0];
1317 36 : for(unsigned core_id = 2; core_id < this->d_; ++core_id) {
1318 14 : B.ref(core_id) = envi_L[core_id - 1] * this->vb_(core_id) * envi_R[core_id - 1];
1319 : }
1320 22 : B.ref(this->d_) = envi_L[this->d_ - 1] * this->vb_(this->d_);
1321 :
1322 22 : return std::make_tuple(B, envi_L, envi_R);
1323 22 : }
1324 :
1325 : #endif // !defined(__PLUMED_HAS_LIBITENSOR) || !defined(__PLUMED_HAS_LIBHDF5)
1326 :
1327 20 : void TTMetaD::registerKeywords(Keywords& keys) {
1328 20 : Bias::registerKeywords(keys);
1329 20 : keys.addFlag("KERNEL_BASIS", false, "Specifies that local kernel basis should be used instead of Fourier basis");
1330 20 : keys.add("compulsory", "SIGMA", "the widths of the Gaussian hills");
1331 20 : keys.add("compulsory", "PACE", "the frequency for hill addition");
1332 20 : keys.add("compulsory", "FILE", "HILLS", "a file in which the list of added hills is stored");
1333 20 : keys.add("compulsory", "TT_FILE", "ttsketch.h5", "the HDF5 file where the TT bias is stored after every sketch, and read at restart. "
1334 : "With multiple replicas the replica suffix is added, with WALKERS_MPI the file is shared and is in the walkers directory");
1335 20 : keys.add("compulsory", "HEIGHT", "the heights of the Gaussian hills");
1336 20 : keys.add("optional", "FMT", "specify format for HILLS files (useful for decrease the number of digits in regtests)");
1337 20 : keys.add("optional", "BIASFACTOR", "use well tempered metadynamics and use this bias factor. Please note you must also specify temp");
1338 20 : keys.add("optional", "TEMP", "the system temperature - this is only needed if you are doing well-tempered metadynamics");
1339 20 : keys.addFlag("WALKERS_MPI", false, "To be used when gromacs + multiple walkers are used");
1340 20 : keys.add("optional", "WALKERS_DIR", "shared directory with the hills files from all the walkers");
1341 20 : keys.use("RESTART");
1342 20 : keys.add("optional", "SKETCH_RANK", "Target rank for TTSketch algorithm - compulsory if SKETCH_CUTOFF is not specified");
1343 20 : keys.add("optional", "SKETCH_CUTOFF", "Truncation error cutoff for singular value decomposition - compulsory if SKETCH_RANK is not specified");
1344 20 : keys.add("compulsory", "SKETCH_INITRANK", "Initial rank for TTSketch algorithm");
1345 20 : keys.add("compulsory", "SKETCH_PACE", "1e6", "The frequency for TT Vbias updates");
1346 20 : keys.add("compulsory", "INTERVAL_MIN", "Lower limits of the domain of the basis functions. Outside the limits the bias is the one at the closest limit and the system does not feel the biasing force along that argument");
1347 20 : keys.add("compulsory", "INTERVAL_MAX", "Upper limits of the domain of the basis functions. Outside the limits the bias is the one at the closest limit and the system does not feel the biasing force along that argument");
1348 20 : keys.add("compulsory", "SKETCH_NBASIS", "20", "Number of basis functions per dimension");
1349 20 : keys.add("compulsory", "SKETCH_ALPHA", "0.05", "Weight coefficient for random tensor train construction");
1350 20 : keys.add("optional", "SKETCH_UNTIL", "After this time, the hills deposited so far are compressed into the TT one last time and the bias potential freezes");
1351 20 : keys.add("optional", "SKETCH_WIDTH", "Width of Gaussian kernels for smoothing");
1352 20 : keys.add("optional", "KERNEL_DX", "Width of basis function kernels");
1353 20 : keys.addFlag("NONINTRUSIVE", false, "Sketching uses previous exact sum of Gaussians instead of TT approximation");
1354 20 : keys.addFlag("DETERMINISTIC", false, "Use a fixed random seed for TT sketch construction, ensuring reproducible results. Intended for regression testing.");
1355 20 : }
1356 :
1357 : PLUMED_REGISTER_ACTION(TTMetaD, "TTMETAD")
1358 :
1359 : } // namespace ttsketch
1360 : } // namespace PLMD
|
