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 "BasisFunc.h"
7 : #include "tools/Exception.h"
8 : #include <limits>
9 :
10 : namespace PLMD {
11 : namespace ttsketch {
12 :
13 0 : BasisFunc::BasisFunc()
14 0 : : dom_(std::make_pair(0.0, 0.0)), nbasis_(0), L_(0.0), shift_(0.0), inv_L_(0.0), sqrt_inv_L_(0.0), inv_sqrt_2L_(0.0), w_(0.0), kernel_(false), dx_(0.0),
15 0 : gramCondition_(1.0), gramDropped_(0) {}
16 :
17 50 : BasisFunc::BasisFunc(std::pair<double, double> dom, int nbasis, double w, bool kernel, double dx)
18 50 : : dom_(dom), nbasis_(nbasis), L_((dom.second - dom.first) * 0.5), shift_((dom.second + dom.first) * 0.5),
19 50 : inv_L_(2.0 / (dom.second - dom.first)), sqrt_inv_L_(sqrt(2.0 / (dom.second - dom.first))),
20 50 : inv_sqrt_2L_(sqrt(1.0 / (dom.second - dom.first))),
21 50 : w_(w), kernel_(kernel), dx_(dx), gramCondition_(1.0), gramDropped_(0) {
22 50 : if(kernel) {
23 9 : double spacing = (dom.second - dom.first) / (nbasis - 1);
24 9 : if(dx == 0) {
25 : // default kernel width equals the grid spacing
26 5 : this->dx_ = spacing;
27 : }
28 : // nbasis-1 kernel centers, uniformly spaced starting at dom.first
29 9 : this->centers_ = std::vector<double>(nbasis - 1);
30 72 : for(int i = 0; i < nbasis - 1; ++i) {
31 63 : this->centers_[i] = dom.first + i * spacing;
32 : }
33 : // Compute the Gram matrix G where G(i,j) = int phi_i(x)*phi_j(x) dx,
34 : // integrating the periodized Gaussian kernels analytically over the domain.
35 : // phi_1 = 1 (constant), phi_{i+1} = sum_{k=-1,0,1} exp(-(x-c_i+2kL)^2/(2*dx^2)).
36 : // The pseudo-inverse of G is stored in ginv_ and applied after sketching to
37 : // recover dual basis coefficients from the Gram representation.
38 9 : this->gram_ = Matrix<double>(nbasis, nbasis);
39 : // G(0,0) = int 1*1 dx = domain length
40 9 : this->gram_(0, 0) = this->dom_.second - this->dom_.first;
41 72 : for(int i = 1; i < nbasis; ++i) {
42 : // G(0,i) = G(i,0) = int phi_{i+1}(x) dx (integral of a periodized Gaussian)
43 63 : this->gram_(i, 0) = this->gram_(0, i) = this->dx_ * sqrt(M_PI * 0.5) *
44 63 : (erf((this->dom_.second -
45 63 : 2 * this->dom_.first + this->centers_[i - 1]) /
46 63 : (sqrt(2) * this->dx_)) -
47 63 : erf((this->dom_.first - 2 * this->dom_.second +
48 : this->centers_[i - 1]) /
49 : (sqrt(2) * this->dx_)));
50 : // G(i,j) = int phi_{i+1}(x)*phi_{j+1}(x) dx; the double sum over k,l
51 : // accounts for cross-terms between periodic images of each kernel.
52 318 : for(int j = i; j < nbasis; ++j) {
53 : double result = 0.0;
54 1020 : for(int k = -1; k <= 1; ++k) {
55 3060 : for(int l = -1; l <= 1; ++l) {
56 2295 : result += this->dx_ * 0.5 * exp(-pow((this->dom_.first -
57 2295 : this->dom_.second) * (k - l) + this->centers_[i - 1] -
58 2295 : this->centers_[j - 1], 2) / (4 * pow(this->dx_, 2))) *
59 2295 : sqrt(M_PI) * (erf((this->dom_.first * (k + l - 2) -
60 2295 : this->dom_.second * (k + l) + this->centers_[i - 1] +
61 2295 : this->centers_[j - 1]) / (2 * this->dx_)) -
62 2295 : erf((this->dom_.first * (k + l) - this->dom_.second *
63 2295 : (k + l + 2) + this->centers_[i - 1] +
64 : this->centers_[j - 1]) / (2 * this->dx_)));
65 : }
66 : }
67 255 : this->gram_(i, j) = this->gram_(j, i) = result;
68 : }
69 : }
70 : // Pseudo-inverse of the (symmetric, positive semi-definite) Gram matrix from its
71 : // eigen-decomposition, ginv = sum_k v_k v_k^T / lambda_k over the eigenvalues
72 : // lambda_k > gramTolerance() * lambda_max. The kernel basis is nearly linearly
73 : // dependent when the kernels are wide compared with their spacing: the constant
74 : // function is then almost a combination of the periodized Gaussians, and the
75 : // Gram matrix has an eigenvalue at the level of the rounding errors (e.g. 1e-14
76 : // relative for dx equal to the spacing). Inverting it would divide by a number
77 : // that is mostly rounding error and amplify it into the dual coefficients, making
78 : // the bias depend on the machine at the percent level.
79 9 : std::vector<double> eval;
80 9 : Matrix<double> evec; // eigenvectors are the rows
81 9 : if(diagMat(this->gram_, eval, evec) != 0) {
82 0 : plumed_merror("diagonalization of the Gram matrix of the kernel basis failed");
83 : }
84 9 : const double lmax = eval[nbasis - 1];
85 9 : const double lmin = eval[0];
86 9 : this->gramCondition_ = lmin > 0.0 ? lmax / lmin : std::numeric_limits<double>::infinity();
87 18 : this->ginv_ = Matrix<double>(nbasis, nbasis);
88 : this->ginv_ = 0.0;
89 81 : for(int k = 0; k < nbasis; ++k) {
90 72 : if(!(eval[k] > gramTolerance() * lmax)) {
91 9 : ++this->gramDropped_;
92 9 : continue;
93 : }
94 63 : const double inv = 1.0 / eval[k];
95 573 : for(int i = 0; i < nbasis; ++i) {
96 4680 : for(int j = 0; j < nbasis; ++j) {
97 4170 : this->ginv_(i, j) += inv * evec(k, i) * evec(k, j);
98 : }
99 : }
100 : }
101 9 : } else {
102 : // damping factors of the harmonics in conv mode (see operator())
103 41 : this->damp_.resize(nbasis / 2 + 1);
104 193 : for(unsigned k = 0; k < this->damp_.size(); ++k) {
105 152 : this->damp_[k] = exp(-0.5 * pow(M_PI * this->w_ * k * this->inv_L_, 2));
106 : }
107 : }
108 50 : }
109 :
110 2347 : void BasisFunc::evalAll(double x, bool conv, double* phi, double* dphi) const {
111 2347 : if(this->kernel_) {
112 : // constant function, then the periodized Gaussians (widened by w_ in conv mode)
113 414 : phi[0] = 1.0;
114 414 : dphi[0] = 0.0;
115 414 : const double s2 = conv ? this->dx_ * this->dx_ + this->w_ * this->w_ : this->dx_ * this->dx_;
116 414 : const double pref = conv ? this->dx_ / sqrt(s2) : 1.0;
117 414 : const double inv_s2 = 1.0 / s2;
118 3384 : for(int pos = 2; pos <= this->nbasis_; ++pos) {
119 : double v = 0.0, g = 0.0;
120 11880 : for(int k = -1; k <= 1; ++k) {
121 8910 : const double u = x - this->centers_[pos - 2] + 2 * k * this->L_;
122 8910 : const double e = pref * exp(-0.5 * u * u * inv_s2);
123 8910 : v += e;
124 8910 : g -= u * inv_s2 * e;
125 : }
126 2970 : phi[pos - 1] = v;
127 2970 : dphi[pos - 1] = g;
128 : }
129 : } else {
130 : // Fourier: phi_{2k} ~ cos(k*theta), phi_{2k+1} ~ sin(k*theta), theta = pi*(x-shift)/L
131 1933 : phi[0] = this->inv_sqrt_2L_;
132 1933 : dphi[0] = 0.0;
133 1933 : const double theta = M_PI * this->inv_L_ * (x - this->shift_);
134 1933 : const double c1 = cos(theta), s1 = sin(theta);
135 : const double dtheta = M_PI * this->inv_L_;
136 : double ck = c1, sk = s1;
137 12883 : for(int pos = 2; pos <= this->nbasis_; ++pos) {
138 10950 : const int k = pos / 2;
139 10950 : if(pos > 2 && pos % 2 == 0) {
140 : // cos(k*theta), sin(k*theta) from k-1
141 3542 : const double cn = ck * c1 - sk * s1;
142 3542 : sk = sk * c1 + ck * s1;
143 : ck = cn;
144 : }
145 10950 : const double a = conv ? this->sqrt_inv_L_ * this->damp_[k] : this->sqrt_inv_L_;
146 10950 : if(pos % 2 == 0) {
147 5475 : phi[pos - 1] = a * ck;
148 5475 : dphi[pos - 1] = -a * k * dtheta * sk;
149 : } else {
150 5475 : phi[pos - 1] = a * sk;
151 5475 : dphi[pos - 1] = a * k * dtheta * ck;
152 : }
153 : }
154 : }
155 2347 : }
156 :
157 : // Orthonormal Fourier basis (not convolved):
158 : // pos=1: phi_1(x) = 1/sqrt(2L)
159 : // pos even: phi_pos(x) = 1/sqrt(L) * cos(pi*(x-shift)*k/L), k = pos/2
160 : // pos odd: phi_pos(x) = 1/sqrt(L) * sin(pi*(x-shift)*k/L), k = pos/2
161 : // k = floor(0.5*pos) intentionally pairs even and odd indices to the same harmonic:
162 : // phi_2 and phi_3 both use k=1, phi_4 and phi_5 use k=2, etc.
163 1326894 : double BasisFunc::fourier(double x, int pos) const {
164 1326894 : if(pos == 1) {
165 120354 : return this->inv_sqrt_2L_;
166 : }
167 1206540 : double k = floor(0.5 * pos);
168 1206540 : double phase = M_PI * k * this->inv_L_ * (x - this->shift_);
169 1206540 : if(pos % 2 == 0) {
170 603270 : return this->sqrt_inv_L_ * cos(phase);
171 : } else {
172 603270 : return this->sqrt_inv_L_ * sin(phase);
173 : }
174 : }
175 :
176 : // Periodized Gaussian kernel basis:
177 : // pos=1: phi_1(x) = 1 (constant)
178 : // pos>1: phi_pos(x) = sum_{k=-1,0,1} exp(-(x - c_{pos-1} + 2kL)^2 / (2*dx^2))
179 : // The k=-1,0,1 sum enforces periodicity across domain boundaries.
180 7565982 : double BasisFunc::gaussian(double x, int pos) const {
181 7565982 : if(pos == 1) {
182 : return 1.0;
183 : } else {
184 : double result = 0.0;
185 26901216 : for(int k = -1; k <= 1; ++k) {
186 20175912 : result += exp(-pow(x - this->centers_[pos - 2] + 2 * k * this->L_, 2) / (2 * pow(this->dx_, 2)));
187 : }
188 : return result;
189 : }
190 : }
191 :
192 4812 : double BasisFunc::fourierd(double x, int pos) const {
193 4812 : if(pos == 1) {
194 : return 0.0;
195 : }
196 : // k is the same harmonic index as in fourier()
197 4560 : double k = floor(0.5 * pos);
198 : // pi*k/L^(3/2) = pi*k * inv_L_ * sqrt_inv_L_, precomputed to avoid pow()
199 4560 : double coeff = M_PI * k * this->inv_L_ * this->sqrt_inv_L_;
200 4560 : double phase = M_PI * k * this->inv_L_ * (x - this->shift_);
201 4560 : if(pos % 2 == 0) {
202 2280 : return -coeff * sin(phase);
203 : } else {
204 2280 : return coeff * cos(phase);
205 : }
206 : }
207 :
208 3876 : double BasisFunc::gaussiand(double x, int pos) const {
209 3876 : if(pos == 1) {
210 : return 0.0;
211 : } else {
212 : double result = 0.0;
213 13728 : for(int k = -1; k <= 1; ++k) {
214 10296 : result += (this->centers_[pos - 2] - x - 2 * k * this->L_) /
215 10296 : pow(this->dx_, 2) * exp(-pow(x - this->centers_[pos - 2] +
216 10296 : 2 * k * this->L_, 2) / (2 * pow(this->dx_, 2)));
217 : }
218 : return result;
219 : }
220 : }
221 :
222 : // Evaluate basis function phi_pos(x), optionally in convolution (conv) mode.
223 : // Conv mode for Fourier: multiplies each harmonic by exp(-(pi*w*k)^2/(2L^2)),
224 : // which is equivalent to convolving the basis function with a Gaussian of width w.
225 : // Conv mode for kernel: widens each kernel analytically to std dev sqrt(dx^2+w^2),
226 : // equivalent to convolving the Gaussian kernel with a Gaussian of width w.
227 : // In both cases the constant basis function (pos=1) is unchanged by convolution.
228 8893224 : double BasisFunc::operator()(double x, int pos, bool conv) const {
229 8893224 : if(this->kernel_) {
230 7565982 : if(conv) {
231 0 : if(pos == 1) {
232 : return 1.0;
233 : } else {
234 : // Gaussian kernel convolved with Gaussian of width w_: std dev becomes sqrt(dx^2+w^2)
235 : double result = 0.0;
236 0 : for(int k = -1; k <= 1; ++k) {
237 0 : result += exp(-pow(x - this->centers_[pos - 2] +
238 0 : 2 * k * this->L_, 2) / (2 * (pow(this->dx_, 2) +
239 0 : pow(this->w_, 2)))) / (sqrt(1 / pow(this->dx_, 2) + 1 /
240 0 : pow(this->w_, 2)) * this->w_);
241 : }
242 : return result;
243 : }
244 : } else {
245 7565982 : return gaussian(x, pos);
246 : }
247 : } else {
248 1327242 : if(conv) {
249 3588 : if(pos == 1) {
250 348 : return this->inv_sqrt_2L_; // constant, unaffected by convolution
251 : } else {
252 : // Fourier convolution theorem: multiplying by exp(-(pi*w*k)^2/(2L^2))
253 : // is equivalent to convolving phi_pos with a Gaussian of width w_
254 3240 : double k = floor(0.5 * pos);
255 3240 : return exp(-0.5 * pow(M_PI * this->w_ * k * this->inv_L_, 2)) * fourier(x, pos);
256 : }
257 : } else {
258 1323654 : return fourier(x, pos);
259 : }
260 : }
261 : }
262 :
263 8940 : double BasisFunc::grad(double x, int pos, bool conv) const {
264 8940 : if(this->kernel_) {
265 3876 : if(conv) {
266 0 : if(pos == 1) {
267 : return 0.0;
268 : } else {
269 : double result = 0.0;
270 0 : for(int k = -1; k <= 1; ++k) {
271 0 : result += pow(this->dx_, 2) * exp(-pow(x - this->centers_[pos - 2] +
272 0 : 2 * k * this->L_, 2) / (2 * (pow(this->dx_, 2) +
273 0 : pow(this->w_, 2)))) * (this->centers_[pos - 2] -
274 0 : 2 * k * this->L_ - x) * sqrt(1 / pow(this->dx_, 2) +
275 0 : 1 / pow(this->w_, 2)) * this->w_ / pow(pow(this->dx_, 2) +
276 0 : pow(this->w_, 2), 2);
277 : }
278 : return result;
279 : }
280 : } else {
281 3876 : return gaussiand(x, pos);
282 : }
283 : } else {
284 5064 : if(conv) {
285 2532 : if(pos == 1) {
286 : return 0.0;
287 : } else {
288 2280 : double k = floor(0.5 * pos);
289 2280 : return exp(-0.5 * pow(M_PI * this->w_ * k * this->inv_L_, 2)) * fourierd(x, pos);
290 : }
291 : } else {
292 2532 : return fourierd(x, pos);
293 : }
294 : }
295 : }
296 :
297 : // Compute int_a^b phi_pos(x) dx analytically.
298 : // Used by covMat() to evaluate the partition function Z = <1> under the TT distribution.
299 40 : double BasisFunc::int0(int pos) const {
300 40 : if(this->kernel_) {
301 18 : if(pos == 1) {
302 2 : return this->dom_.second - this->dom_.first;
303 : } else {
304 16 : return -this->dx_ * sqrt(M_PI * 0.5) * (erf((2 * this->dom_.first -
305 16 : this->dom_.second - this->centers_[pos - 2]) /
306 16 : (sqrt(2) * this->dx_)) + erf((this->dom_.first -
307 16 : 2 * this->dom_.second + this->centers_[pos - 2]) /
308 16 : (sqrt(2) * this->dx_)));
309 : }
310 : } else {
311 22 : if(pos == 1) {
312 2 : return sqrt(2 * this->L_);
313 20 : } else if(pos % 2 == 0) {
314 10 : double k = floor(0.5 * pos);
315 10 : return 2 * sqrt(this->L_) / (M_PI * k) * sin(M_PI * k);
316 : } else {
317 : return 0.0;
318 : }
319 : }
320 : }
321 :
322 : // Compute int_a^b x*phi_pos(x) dx analytically.
323 : // Used by covMat() to evaluate E[x] = (1/Z) * <x> under the TT distribution.
324 40 : double BasisFunc::int1(int pos) const {
325 40 : if(this->kernel_) {
326 18 : if(pos == 1) {
327 2 : return (pow(this->dom_.second, 2) - pow(this->dom_.first, 2)) * 0.5;
328 : } else {
329 16 : double a = this->dom_.first;
330 16 : double b = this->dom_.second;
331 16 : double c = this->centers_[pos - 2];
332 16 : double dx_sq = pow(this->dx_, 2);
333 : double sqrt2pi = sqrt(2 * M_PI);
334 16 : double term1 = exp(-pow(b - 2 * a + c, 2) / (2 * dx_sq)) -
335 16 : exp(-pow(a - 2 * b + c, 2) / (2 * dx_sq));
336 16 : double term2 = (b - a + c) * sqrt2pi * erf((a - c) / (sqrt(2) * this->dx_));
337 16 : double term3 = (a - b - c) * sqrt2pi * erf((2 * a - b - c) / (sqrt(2) * this->dx_));
338 16 : double term4 = sqrt2pi * (c * erf((-a + c) / (sqrt(2) * this->dx_)) -
339 16 : (a - b + c) * erf((a - 2 * b + c) / (sqrt(2) * this->dx_)) +
340 16 : (a - b) * erf((-b + c) / (sqrt(2) * this->dx_)));
341 16 : return this->dx_ * 0.5 * (2 * this->dx_ * term1 + term2 + term3 + term4);
342 : }
343 : } else {
344 22 : if(pos == 1) {
345 2 : return this->shift_ * sqrt(2 * this->L_);
346 20 : } else if(pos % 2 == 0) {
347 10 : double k = floor(0.5 * pos);
348 10 : return 2 * this->shift_ * sqrt(this->L_) / (M_PI * k) * sin(M_PI * k);
349 : } else {
350 10 : double k = floor(0.5 * pos);
351 10 : return 2 * pow(this->L_, 1.5) / pow(M_PI * k, 2) * (sin(M_PI * k) - M_PI * k * cos(M_PI * k));
352 : }
353 : }
354 : }
355 :
356 : // Compute int_a^b x^2*phi_pos(x) dx analytically.
357 : // Used by covMat() to evaluate E[x^2] for variance: Var[x] = E[x^2] - E[x]^2.
358 40 : double BasisFunc::int2(int pos) const {
359 40 : if(this->kernel_) {
360 18 : if(pos == 1) {
361 2 : return (pow(this->dom_.second, 3) - pow(this->dom_.first, 3)) / 3;
362 : } else {
363 16 : double a = this->dom_.first;
364 16 : double b = this->dom_.second;
365 16 : double c = this->centers_[pos - 2];
366 16 : double dx_sq = pow(this->dx_, 2);
367 : double sqrt2pi = sqrt(2 * M_PI);
368 16 : double sqrt2_dx = sqrt(2) * this->dx_;
369 16 : double exp1 = exp(-pow(a - c, 2) / (2 * dx_sq));
370 16 : double exp2 = exp(-pow(b - c, 2) / (2 * dx_sq));
371 16 : double exp3 = exp(-pow(a - 2 * b + c, 2) / (2 * dx_sq));
372 16 : double exp4 = exp(-pow(b - 2 * a + c, 2) / (2 * dx_sq));
373 16 : double term1 = 2 * this->dx_ * (
374 16 : a * (2 * exp1 + 2 * exp2 - exp3) +
375 16 : b * (exp4 - 2 * exp1 - 2 * exp2) +
376 16 : c * (exp4 - exp3)
377 : );
378 16 : double diff1_sq = pow(b - a + c, 2) + dx_sq;
379 16 : double diff2_sq = pow(a - b + c, 2) + dx_sq;
380 16 : double c_sq = pow(c, 2) + dx_sq;
381 16 : double erf1 = erf((a - c) / sqrt2_dx);
382 16 : double erf2 = erf((2 * a - b - c) / sqrt2_dx);
383 16 : double erf3 = erf((c - a) / sqrt2_dx);
384 16 : double erf4 = erf((a - 2 * b + c) / sqrt2_dx);
385 16 : double erf5 = erf((c - b) / sqrt2_dx);
386 16 : double term2 = diff1_sq * sqrt2pi * erf1 - diff1_sq * sqrt2pi * erf2;
387 16 : double term3 = sqrt2pi * (
388 16 : c_sq * erf3 - diff2_sq * erf4 +
389 16 : (a - b) * (a - b + 2 * c) * erf5
390 : );
391 16 : return this->dx_ * 0.5 * (term1 + term2 + term3);
392 : }
393 : } else {
394 22 : if(pos == 1) {
395 2 : return sqrt(2 * this->L_) / 3 * (3 * pow(this->shift_, 2) + pow(this->L_, 2));
396 20 : } else if(pos % 2 == 0) {
397 10 : double k = floor(0.5 * pos);
398 10 : return 2 * sqrt(this->L_) / pow(M_PI * k, 3) *
399 10 : (2 * pow(this->L_, 2) * M_PI * k * cos(M_PI * k) +
400 10 : (pow(M_PI * k, 2) * (pow(this->shift_, 2) +
401 10 : pow(this->L_, 2)) - 2 * pow(this->L_, 2)) * sin(M_PI * k));
402 : } else {
403 10 : double k = floor(0.5 * pos);
404 10 : return 4 * this->shift_ * pow(this->L_, 1.5) / pow(M_PI * k, 2) * (sin(M_PI * k) - M_PI * k * cos(M_PI * k));
405 : }
406 : }
407 : }
408 :
409 : }
410 : }
|
