(moose/modules/stochastic_tools/src/surrogates/GaussianProcessSurrogate.C)
// This file is part of the MOOSE framework
// https://mooseframework.inl.gov
//
// All rights reserved, see COPYRIGHT for full restrictions
// https://github.com/idaholab/moose/blob/master/COPYRIGHT
//
// Licensed under LGPL 2.1, please see LICENSE for details
// https://www.gnu.org/licenses/lgpl-2.1.html
#ifdef MOOSE_LIBTORCH_ENABLED
#include "GaussianProcessSurrogate.h"
#include "Sampler.h"
#include "CovarianceFunctionBase.h"
registerMooseObject("StochasticToolsApp", GaussianProcessSurrogate);
InputParameters
GaussianProcessSurrogate::validParams()
{
InputParameters params = SurrogateModel::validParams();
params.addClassDescription("Computes and evaluates Gaussian Process surrogate model.");
return params;
}
GaussianProcessSurrogate::GaussianProcessSurrogate(const InputParameters & parameters)
: SurrogateModel(parameters),
CovarianceInterface(parameters),
_gp(declareModelData<StochasticTools::GaussianProcess>("_gp")),
_training_params(getModelData<torch::Tensor>("_training_params"))
{
}
void
GaussianProcessSurrogate::setupCovariance(UserObjectName covar_name)
{
if (_gp.getCovarFunctionPtr() != nullptr)
::mooseError("Attempting to redefine covariance function using setupCovariance.");
_gp.linkCovarianceFunction(getCovarianceFunctionByName(covar_name));
}
Real
GaussianProcessSurrogate::evaluate(const std::vector<Real> & x) const
{
// Overlaod for evaluate to maintain general compatibility. Only returns mean
Real dummy = 0;
return this->evaluate(x, dummy);
}
Real
GaussianProcessSurrogate::evaluate(const std::vector<Real> & x, Real & std_dev) const
{
std::vector<Real> y;
std::vector<Real> std;
this->evaluate(x, y, std);
std_dev = std[0];
return y[0];
}
void
GaussianProcessSurrogate::evaluate(const std::vector<Real> & x, std::vector<Real> & y) const
{
// Overlaod for evaluate to maintain general compatibility. Only returns mean
std::vector<Real> std_dummy;
this->evaluate(x, y, std_dummy);
}
void
GaussianProcessSurrogate::evaluate(const std::vector<Real> & x,
std::vector<Real> & y,
std::vector<Real> & std) const
{
const unsigned int n_dims = _training_params.sizes()[1];
mooseAssert(x.size() == n_dims,
"Number of parameters provided for evaluation does not match number of parameters "
"used for training.");
const unsigned int n_outputs = _gp.getCovarFunction().numOutputs();
y = std::vector<Real>(n_outputs, 0.0);
std = std::vector<Real>(n_outputs, 0.0);
const auto options = _training_params.options().dtype(at::kDouble);
torch::Tensor test_points = torch::empty({1, n_dims}, at::kDouble);
auto points_accessor = test_points.accessor<Real, 2>();
for (unsigned int ii = 0; ii < n_dims; ++ii)
points_accessor[0][ii] = x[ii];
test_points = test_points.to(options.device());
_gp.getParamStandardizer().getStandardized(test_points);
torch::Tensor K_train_test =
torch::empty({_training_params.sizes()[0] * n_outputs, n_outputs}, options);
_gp.getCovarFunction().computeCovarianceMatrix(
K_train_test, _training_params, test_points, false);
torch::Tensor K_test = torch::empty({n_outputs, n_outputs}, options);
_gp.getCovarFunction().computeCovarianceMatrix(K_test, test_points, test_points, true);
// Compute the predicted mean value (centered)
torch::Tensor pred_value = torch::transpose(
torch::mm(torch::transpose(K_train_test, 0, 1), _gp.getKResultsSolve()), 0, 1);
// De-center/scale the value and store for return
_gp.getDataStandardizer().getDestandardized(pred_value);
torch::Tensor pred_var =
K_test - torch::mm(torch::transpose(K_train_test, 0, 1),
torch::cholesky_solve(K_train_test, _gp.getKCholeskyDecomp()));
// Only the marginal variances are returned. Clamp tiny negative roundoff before sqrt.
torch::Tensor std_dev_vec =
torch::sqrt(torch::clamp_min(torch::diagonal(pred_var), 0.0)).unsqueeze(0);
_gp.getDataStandardizer().getDescaled(std_dev_vec);
const auto std_dev_cpu = LibtorchUtils::toCPUContiguous(std_dev_vec);
const auto pred_value_cpu = LibtorchUtils::toCPUContiguous(pred_value);
auto std_accessor = std_dev_cpu.accessor<Real, 2>();
auto pred_value_accessor = pred_value_cpu.accessor<Real, 2>();
for (const auto output_i : make_range(n_outputs))
{
y[output_i] = pred_value_accessor[0][output_i];
std[output_i] = std_accessor[0][output_i];
}
}
#endif