Line data Source code
1 : //* This file is part of the MOOSE framework
2 : //* https://mooseframework.inl.gov
3 : //*
4 : //* All rights reserved, see COPYRIGHT for full restrictions
5 : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
6 : //*
7 : //* Licensed under LGPL 2.1, please see LICENSE for details
8 : //* https://www.gnu.org/licenses/lgpl-2.1.html
9 : #ifdef MOOSE_LIBTORCH_ENABLED
10 :
11 : #include "GaussianProcess.h"
12 : #include "FEProblemBase.h"
13 :
14 : #include <petsctao.h>
15 : #include <petscdmda.h>
16 :
17 : #include "libmesh/petsc_vector.h"
18 : #include "libmesh/petsc_matrix.h"
19 :
20 : #include <cmath>
21 :
22 : #include "MooseRandom.h"
23 : #include "Shuffle.h"
24 :
25 : #include <torch/optim/adam.h>
26 :
27 : namespace StochasticTools
28 : {
29 :
30 : namespace
31 : {
32 :
33 : using HyperParameterMap = GaussianProcess::HyperParameterMap;
34 :
35 : torch::Tensor
36 368403 : flattenOutputData(const torch::Tensor & output_data)
37 : {
38 : mooseAssert(output_data.dim() == 2, "GaussianProcess output data must be rank-2.");
39 368403 : return torch::reshape(torch::transpose(output_data, 0, 1),
40 736806 : {output_data.size(0) * output_data.size(1), 1});
41 : }
42 :
43 : torch::TensorOptions
44 1120694 : doubleOptionsLike(const torch::Tensor & tensor)
45 : {
46 1120694 : return tensor.options().dtype(at::kDouble);
47 : }
48 :
49 : torch::Tensor
50 754087 : toOptions(const torch::Tensor & tensor, const torch::TensorOptions & options)
51 : {
52 754087 : auto result = tensor.to(options.device());
53 754087 : if (result.scalar_type() != at::kDouble)
54 0 : result = result.to(at::kDouble);
55 754087 : return result;
56 : }
57 :
58 : std::vector<Real>
59 598 : exportHyperParameter(const torch::Tensor & tensor)
60 : {
61 905 : if (!CovarianceFunctionBase::isScalarHyperParameter(tensor) &&
62 307 : !CovarianceFunctionBase::isVectorHyperParameter(tensor))
63 0 : mooseError("Unsupported hyperparameter rank ", tensor.dim(), ".");
64 598 : auto cpu_tensor = LibtorchUtils::toCPUContiguous(tensor);
65 598 : if (cpu_tensor.scalar_type() != at::kDouble)
66 0 : cpu_tensor = cpu_tensor.to(at::kDouble).contiguous();
67 598 : const auto flattened = cpu_tensor.reshape({-1});
68 1794 : return {flattened.data_ptr<Real>(), flattened.data_ptr<Real>() + flattened.numel()};
69 : }
70 :
71 : torch::Tensor
72 384000 : buildVectorHyperParameter(const std::vector<Real> & values, const torch::TensorOptions & options)
73 : {
74 384000 : auto tensor = torch::empty({long(values.size())}, torch::TensorOptions().dtype(at::kDouble));
75 384000 : auto tensor_accessor = tensor.accessor<Real, 1>();
76 1197000 : for (const auto index : index_range(values))
77 813000 : tensor_accessor[index] = values[index];
78 768000 : return toOptions(tensor, options);
79 : }
80 :
81 : void
82 403 : moveHyperParameters(HyperParameterMap & hyperparameters, const torch::TensorOptions & options)
83 : {
84 1684 : for (auto & iter : hyperparameters)
85 2562 : iter.second = toOptions(iter.second, options);
86 403 : }
87 :
88 : void
89 752000 : updateHyperParameter(torch::Tensor & tensor,
90 : const std::vector<Real> & values,
91 : const std::string & name)
92 : {
93 752000 : const auto options = doubleOptionsLike(tensor);
94 752000 : if (CovarianceFunctionBase::isScalarHyperParameter(tensor))
95 : {
96 : mooseAssert(values.size() == 1, "Scalar hyperparameter update requires a single value.");
97 : tensor =
98 736000 : toOptions(torch::tensor(values[0], torch::TensorOptions().dtype(at::kDouble)), options);
99 : }
100 384000 : else if (CovarianceFunctionBase::isVectorHyperParameter(tensor))
101 768000 : tensor = buildVectorHyperParameter(values, options);
102 : else
103 0 : mooseError("Unsupported hyperparameter rank ", tensor.dim(), " for ", name, ".");
104 752000 : }
105 :
106 : } // namespace
107 :
108 226 : GaussianProcess::GPOptimizerOptions::GPOptimizerOptions(const unsigned int show_every_nth_iteration,
109 : const unsigned int num_iter,
110 : const unsigned int batch_size,
111 : const Real learning_rate,
112 : const Real b1,
113 : const Real b2,
114 : const Real eps,
115 : const Real lambda,
116 226 : const OptimizerType optimizer_type)
117 226 : : show_every_nth_iteration(show_every_nth_iteration),
118 226 : num_iter(num_iter),
119 226 : batch_size(batch_size),
120 226 : learning_rate(learning_rate),
121 226 : b1(b1),
122 226 : b2(b2),
123 226 : eps(eps),
124 226 : lambda(lambda),
125 226 : optimizer_type(optimizer_type)
126 : {
127 226 : }
128 :
129 458 : GaussianProcess::GaussianProcess() {}
130 :
131 : void
132 224 : GaussianProcess::initialize(CovarianceFunctionBase * covariance_function,
133 : const std::vector<std::string> & params_to_tune,
134 : const std::vector<Real> & min,
135 : const std::vector<Real> & max)
136 : {
137 224 : linkCovarianceFunction(covariance_function);
138 224 : generateTuningMap(params_to_tune, min, max);
139 224 : }
140 :
141 : void
142 240 : GaussianProcess::linkCovarianceFunction(CovarianceFunctionBase * covariance_function)
143 : {
144 240 : _covariance_function = covariance_function;
145 240 : _covar_type = _covariance_function->type();
146 240 : _covar_name = _covariance_function->name();
147 240 : _covariance_function->dependentCovarianceTypes(_dependent_covar_types);
148 240 : _dependent_covar_names = _covariance_function->dependentCovarianceNames();
149 240 : _num_outputs = _covariance_function->numOutputs();
150 240 : }
151 :
152 : void
153 403 : GaussianProcess::setupCovarianceMatrix(const torch::Tensor & training_params,
154 : const torch::Tensor & training_data,
155 : const GPOptimizerOptions & opts)
156 : {
157 403 : const auto options = doubleOptionsLike(training_params);
158 403 : const auto params = toOptions(training_params, options);
159 403 : const auto data = toOptions(training_data, options);
160 :
161 : mooseAssert(params.dim() == 2, "GaussianProcess training parameters must be rank-2.");
162 : mooseAssert(data.dim() == 2, "GaussianProcess training responses must be rank-2.");
163 :
164 : const auto num_samples = params.size(0);
165 : mooseAssert(data.size(0) == num_samples,
166 : "Training parameter and response sample counts must match.");
167 : mooseAssert(data.size(1) == _num_outputs,
168 : "Training response dimension does not match the covariance output dimension.");
169 :
170 403 : const bool batch_decision = opts.batch_size > 0 && (opts.batch_size <= num_samples);
171 403 : _batch_size = batch_decision ? opts.batch_size : num_samples;
172 :
173 : _hyperparam_map.clear();
174 403 : _covariance_function->buildHyperParamMap(_hyperparam_map);
175 403 : moveHyperParameters(_hyperparam_map, options);
176 403 : _covariance_function->loadHyperParamMap(_hyperparam_map);
177 :
178 403 : if (_tuning_data.size())
179 291 : tuneHyperParamsAdam(params, data, opts);
180 :
181 403 : _covariance_function->computeCovarianceMatrix(_K, params, params, true);
182 403 : const auto flattened_tensor = flattenOutputData(data);
183 :
184 : // Compute the Cholesky decomposition and inverse action of the covariance matrix.
185 403 : setupStoredMatrices(flattened_tensor);
186 403 : }
187 :
188 : void
189 368403 : GaussianProcess::setupStoredMatrices(const torch::Tensor & input)
190 : {
191 368403 : _K_cho_decomp = torch::linalg_cholesky(_K);
192 368403 : _K_results_solve = torch::cholesky_solve(input, _K_cho_decomp);
193 368403 : }
194 :
195 : void
196 224 : GaussianProcess::generateTuningMap(const std::vector<std::string> & params_to_tune,
197 : const std::vector<Real> & min_vector,
198 : const std::vector<Real> & max_vector)
199 : {
200 224 : _num_tunable = 0;
201 :
202 : const bool upper_bounds_specified = min_vector.size();
203 : const bool lower_bounds_specified = max_vector.size();
204 :
205 560 : for (const auto param_i : index_range(params_to_tune))
206 : {
207 : const auto & hp = params_to_tune[param_i];
208 336 : if (_covariance_function->isTunable(hp))
209 : {
210 : unsigned int size;
211 : Real min;
212 : Real max;
213 : // Get size and default min/max
214 336 : const bool found = _covariance_function->getTuningData(hp, size, min, max);
215 :
216 336 : if (!found)
217 0 : ::mooseError("The covariance parameter ", hp, " could not be found!");
218 :
219 : // Check for overridden min/max
220 336 : min = lower_bounds_specified ? min_vector[param_i] : min;
221 336 : max = upper_bounds_specified ? max_vector[param_i] : max;
222 : // Save data in tuple
223 : _tuning_data[hp] = std::make_tuple(_num_tunable, size, min, max);
224 336 : _num_tunable += size;
225 : }
226 : }
227 224 : }
228 :
229 : void
230 403 : GaussianProcess::standardizeParameters(torch::Tensor & data, bool keep_moments)
231 : {
232 403 : if (!keep_moments)
233 403 : _param_standardizer.computeSet(data);
234 403 : _param_standardizer.getStandardized(data);
235 403 : }
236 :
237 : void
238 403 : GaussianProcess::standardizeData(torch::Tensor & data, bool keep_moments)
239 : {
240 403 : if (!keep_moments)
241 403 : _data_standardizer.computeSet(data);
242 403 : _data_standardizer.getStandardized(data);
243 403 : }
244 :
245 : void
246 291 : GaussianProcess::tuneHyperParamsAdam(const torch::Tensor & training_params,
247 : const torch::Tensor & training_data,
248 : const GPOptimizerOptions & opts)
249 : {
250 291 : const auto options = doubleOptionsLike(training_params);
251 291 : std::vector<Real> theta_values(_num_tunable, 0.0);
252 :
253 291 : mapToVec(_tuning_data, _hyperparam_map, theta_values);
254 :
255 0 : auto theta = torch::from_blob(theta_values.data(),
256 291 : {static_cast<long>(_num_tunable)},
257 291 : torch::TensorOptions().dtype(at::kDouble))
258 291 : .clone()
259 291 : .to(options.device());
260 :
261 291 : auto adam_options = torch::optim::AdamOptions(opts.learning_rate);
262 : adam_options.betas(std::make_tuple(opts.b1, opts.b2));
263 : adam_options.eps(opts.eps);
264 : // The legacy MOOSE shrink term is decoupled and not learning-rate-scaled, so it cannot be
265 : // represented by Adam's coupled weight_decay option.
266 : adam_options.weight_decay(0.0);
267 582 : torch::optim::Adam optimizer({theta}, adam_options);
268 :
269 : Real store_loss = 0.0;
270 : std::vector<Real> grad_values;
271 291 : const bool use_legacy_update = opts.optimizer_type == OptimizerType::LegacyAdam;
272 :
273 291 : const bool use_full_batch = _batch_size == static_cast<unsigned int>(training_params.size(0));
274 : // Preserve the existing deterministic shuffle sequence for mini-batches, but avoid rebuilding
275 : // shuffled full-batch tensors when the batch already contains every training sample.
276 : std::vector<unsigned int> v_sequence;
277 291 : if (!use_full_batch)
278 : {
279 8 : v_sequence.resize(training_params.size(0));
280 : std::iota(std::begin(v_sequence), std::end(v_sequence), 0);
281 : }
282 291 : if (opts.show_every_nth_iteration)
283 : Moose::out << "OPTIMIZING GP HYPER-PARAMETERS USING "
284 16 : << (use_legacy_update ? "legacy-compatible Adam" : "Adam") << std::endl;
285 368291 : for (unsigned int ss = 0; ss < opts.num_iter; ++ss)
286 : {
287 : torch::Tensor inputs;
288 : torch::Tensor outputs;
289 368000 : if (use_full_batch)
290 : {
291 : inputs = training_params;
292 : outputs = training_data;
293 : }
294 : else
295 : {
296 : MooseRandom generator;
297 8000 : generator.seed(0, 1980);
298 8000 : generator.saveState();
299 : MooseUtils::shuffle<unsigned int>(v_sequence, generator, 0);
300 :
301 8000 : std::vector<int64_t> batch_indices_vec(v_sequence.begin(), v_sequence.begin() + _batch_size);
302 : auto batch_indices = torch::tensor(
303 8000 : batch_indices_vec, torch::TensorOptions().dtype(torch::kLong).device(options.device()));
304 8000 : inputs = torch::index_select(training_params, 0, batch_indices);
305 8000 : outputs = torch::index_select(training_data, 0, batch_indices);
306 8000 : }
307 :
308 368000 : store_loss = getLoss(inputs, outputs);
309 368000 : if (opts.show_every_nth_iteration && ((ss + 1) % opts.show_every_nth_iteration == 0))
310 : Moose::out << "Iteration: " << ss + 1 << " LOSS: " << store_loss << std::endl;
311 :
312 736000 : grad_values = getGradient(inputs);
313 0 : auto grad = torch::from_blob(grad_values.data(),
314 368000 : {static_cast<long>(_num_tunable)},
315 368000 : torch::TensorOptions().dtype(at::kDouble))
316 368000 : .clone()
317 368000 : .to(options.device());
318 368000 : optimizer.zero_grad();
319 : theta.mutable_grad() = grad;
320 : torch::Tensor theta_before_step;
321 368000 : if (use_legacy_update)
322 16000 : theta_before_step = theta.detach().clone();
323 368000 : optimizer.step();
324 :
325 : {
326 368000 : torch::NoGradGuard no_grad;
327 368000 : if (use_legacy_update)
328 24000 : theta -= opts.lambda * theta_before_step;
329 1120000 : for (auto iter = _tuning_data.begin(); iter != _tuning_data.end(); ++iter)
330 : {
331 752000 : const auto first_index = std::get<0>(iter->second);
332 752000 : const auto num_entries = std::get<1>(iter->second);
333 1504000 : const auto min_value = std::get<2>(iter->second);
334 752000 : const auto max_value = std::get<3>(iter->second);
335 2256000 : theta.slice(0, first_index, first_index + num_entries).clamp_(min_value, max_value);
336 : }
337 : }
338 :
339 368000 : const auto theta_export = LibtorchUtils::toCPUContiguous(theta);
340 368000 : const auto * theta_data = theta_export.data_ptr<Real>();
341 368000 : theta_values.assign(theta_data, theta_data + theta_export.numel());
342 368000 : vecToMap(_tuning_data, _hyperparam_map, theta_values);
343 368000 : _covariance_function->loadHyperParamMap(_hyperparam_map);
344 : }
345 291 : if (opts.show_every_nth_iteration)
346 : {
347 : Moose::out << "OPTIMIZED GP HYPER-PARAMETERS:" << std::endl;
348 32 : Moose::out << Moose::stringify(theta_values) << std::endl;
349 : Moose::out << "FINAL LOSS: " << store_loss << std::endl;
350 : }
351 :
352 291 : if (theta_values.size() > 0)
353 : {
354 : unsigned int count = 1;
355 291 : _length_scales.resize(_num_tunable - count);
356 1022 : for (unsigned int i = 0; i < _num_tunable - count; ++i)
357 731 : _length_scales[i] = theta_values[i + 1];
358 : }
359 582 : }
360 :
361 : Real
362 368000 : GaussianProcess::getLoss(torch::Tensor & inputs, torch::Tensor & outputs)
363 : {
364 368000 : _covariance_function->computeCovarianceMatrix(_K, inputs, inputs, true);
365 368000 : const auto flattened_data = flattenOutputData(outputs);
366 :
367 368000 : setupStoredMatrices(flattened_data);
368 :
369 : Real log_likelihood = 0;
370 : log_likelihood +=
371 736000 : -1 * torch::mm(torch::transpose(flattened_data, 0, 1), _K_results_solve).item<Real>();
372 1104000 : log_likelihood += -2.0 * torch::sum(torch::log(torch::diagonal(_K_cho_decomp))).item<Real>();
373 368000 : log_likelihood -= flattened_data.size(0) * std::log(2 * M_PI);
374 368000 : log_likelihood = -log_likelihood / 2;
375 368000 : return log_likelihood;
376 : }
377 :
378 : std::vector<Real>
379 368000 : GaussianProcess::getGradient(torch::Tensor & inputs) const
380 : {
381 368000 : torch::Tensor dKdhp = torch::empty({_num_outputs * _batch_size, _num_outputs * _batch_size},
382 368000 : doubleOptionsLike(inputs));
383 : std::vector<Real> grad_vec;
384 368000 : grad_vec.resize(_num_tunable);
385 1120000 : for (auto iter = _tuning_data.begin(); iter != _tuning_data.end(); ++iter)
386 : {
387 752000 : std::string hyper_param_name = iter->first;
388 752000 : const auto first_index = std::get<0>(iter->second);
389 752000 : const auto num_entries = std::get<1>(iter->second);
390 1933000 : for (unsigned int ii = 0; ii < num_entries; ++ii)
391 : {
392 1181000 : const auto global_index = first_index + ii;
393 1181000 : _covariance_function->computedKdhyper(dKdhp, inputs, hyper_param_name, ii);
394 : const auto quadratic_form =
395 2362000 : torch::mm(torch::transpose(_K_results_solve, 0, 1), torch::mm(dKdhp, _K_results_solve))
396 1181000 : .item<Real>();
397 : const auto inverse_trace =
398 2362000 : torch::trace(torch::cholesky_solve(dKdhp, _K_cho_decomp)).item<Real>();
399 1181000 : grad_vec[global_index] = (inverse_trace - quadratic_form) / 2.0;
400 : }
401 : }
402 368000 : return grad_vec;
403 0 : }
404 :
405 : void
406 291 : GaussianProcess::mapToVec(
407 : const std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> &
408 : tuning_data,
409 : const HyperParameterMap & hyperparam_map,
410 : std::vector<Real> & vec) const
411 : {
412 889 : for (auto iter : tuning_data)
413 : {
414 : const std::string & param_name = iter.first;
415 : const auto tensor_it = hyperparam_map.find(param_name);
416 598 : if (tensor_it == hyperparam_map.end())
417 0 : mooseError("The covariance parameter ", param_name, " could not be found!");
418 :
419 598 : const auto values = exportHyperParameter(tensor_it->second);
420 598 : const auto num_entries = std::get<1>(iter.second);
421 : mooseAssert(values.size() == num_entries,
422 : "Hyperparameter size does not match tuning metadata.");
423 1620 : for (unsigned int ii = 0; ii < num_entries; ++ii)
424 1022 : vec[std::get<0>(iter.second) + ii] = values[ii];
425 598 : }
426 291 : }
427 :
428 : void
429 368000 : GaussianProcess::vecToMap(
430 : const std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> &
431 : tuning_data,
432 : HyperParameterMap & hyperparam_map,
433 : const std::vector<Real> & vec) const
434 : {
435 1120000 : for (auto iter : tuning_data)
436 : {
437 : const std::string & param_name = iter.first;
438 : const auto tensor_it = hyperparam_map.find(param_name);
439 752000 : if (tensor_it == hyperparam_map.end())
440 0 : mooseError("The covariance parameter ", param_name, " could not be found!");
441 :
442 752000 : const auto first_index = std::get<0>(iter.second);
443 752000 : const auto num_entries = std::get<1>(iter.second);
444 752000 : std::vector<Real> values(num_entries);
445 1933000 : for (unsigned int ii = 0; ii < num_entries; ++ii)
446 1181000 : values[ii] = vec[first_index + ii];
447 :
448 752000 : updateHyperParameter(tensor_it->second, values, param_name);
449 752000 : }
450 368000 : }
451 :
452 : } // StochasticTools namespace
453 :
454 : template <>
455 : void
456 23 : dataStore(std::ostream & stream, StochasticTools::GaussianProcess & gp_utils, void * context)
457 : {
458 23 : dataStore(stream, gp_utils.hyperparamMap(), context);
459 23 : dataStore(stream, gp_utils.covarType(), context);
460 23 : dataStore(stream, gp_utils.covarName(), context);
461 : dataStore(stream, gp_utils.covarNumOutputs(), context);
462 23 : dataStore(stream, gp_utils.dependentCovarNames(), context);
463 23 : dataStore(stream, gp_utils.dependentCovarTypes(), context);
464 23 : dataStore(stream, gp_utils.K(), context);
465 23 : dataStore(stream, gp_utils.KResultsSolve(), context);
466 23 : dataStore(stream, gp_utils.KCholeskyDecomp(), context);
467 23 : dataStore(stream, gp_utils.paramStandardizer(), context);
468 23 : dataStore(stream, gp_utils.dataStandardizer(), context);
469 23 : }
470 :
471 : template <>
472 : void
473 16 : dataLoad(std::istream & stream, StochasticTools::GaussianProcess & gp_utils, void * context)
474 : {
475 16 : dataLoad(stream, gp_utils.hyperparamMap(), context);
476 16 : dataLoad(stream, gp_utils.covarType(), context);
477 16 : dataLoad(stream, gp_utils.covarName(), context);
478 : dataLoad(stream, gp_utils.covarNumOutputs(), context);
479 16 : dataLoad(stream, gp_utils.dependentCovarNames(), context);
480 16 : dataLoad(stream, gp_utils.dependentCovarTypes(), context);
481 16 : dataLoad(stream, gp_utils.K(), context);
482 16 : dataLoad(stream, gp_utils.KResultsSolve(), context);
483 16 : dataLoad(stream, gp_utils.KCholeskyDecomp(), context);
484 16 : dataLoad(stream, gp_utils.paramStandardizer(), context);
485 16 : dataLoad(stream, gp_utils.dataStandardizer(), context);
486 16 : }
487 :
488 : #endif
|
