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 "GaussianProcessSurrogate.h"
      12             : #include "Sampler.h"
      13             : 
      14             : #include "CovarianceFunctionBase.h"
      15             : 
      16             : registerMooseObject("StochasticToolsApp", GaussianProcessSurrogate);
      17             : 
      18             : InputParameters
      19         466 : GaussianProcessSurrogate::validParams()
      20             : {
      21         466 :   InputParameters params = SurrogateModel::validParams();
      22         466 :   params.addClassDescription("Computes and evaluates Gaussian Process surrogate model.");
      23         466 :   return params;
      24           0 : }
      25             : 
      26         232 : GaussianProcessSurrogate::GaussianProcessSurrogate(const InputParameters & parameters)
      27             :   : SurrogateModel(parameters),
      28             :     CovarianceInterface(parameters),
      29         232 :     _gp(declareModelData<StochasticTools::GaussianProcess>("_gp")),
      30         696 :     _training_params(getModelData<torch::Tensor>("_training_params"))
      31             : {
      32         232 : }
      33             : 
      34             : void
      35          16 : GaussianProcessSurrogate::setupCovariance(UserObjectName covar_name)
      36             : {
      37          16 :   if (_gp.getCovarFunctionPtr() != nullptr)
      38           0 :     ::mooseError("Attempting to redefine covariance function using setupCovariance.");
      39          16 :   _gp.linkCovarianceFunction(getCovarianceFunctionByName(covar_name));
      40          16 : }
      41             : 
      42             : Real
      43         375 : GaussianProcessSurrogate::evaluate(const std::vector<Real> & x) const
      44             : {
      45             :   // Overlaod for evaluate to maintain general compatibility. Only returns mean
      46         375 :   Real dummy = 0;
      47         375 :   return this->evaluate(x, dummy);
      48             : }
      49             : 
      50             : Real
      51      133926 : GaussianProcessSurrogate::evaluate(const std::vector<Real> & x, Real & std_dev) const
      52             : {
      53             :   std::vector<Real> y;
      54             :   std::vector<Real> std;
      55      133926 :   this->evaluate(x, y, std);
      56      133926 :   std_dev = std[0];
      57      133926 :   return y[0];
      58      133926 : }
      59             : 
      60             : void
      61           0 : GaussianProcessSurrogate::evaluate(const std::vector<Real> & x, std::vector<Real> & y) const
      62             : {
      63             :   // Overlaod for evaluate to maintain general compatibility. Only returns mean
      64             :   std::vector<Real> std_dummy;
      65           0 :   this->evaluate(x, y, std_dummy);
      66           0 : }
      67             : 
      68             : void
      69      134001 : GaussianProcessSurrogate::evaluate(const std::vector<Real> & x,
      70             :                                    std::vector<Real> & y,
      71             :                                    std::vector<Real> & std) const
      72             : {
      73      134001 :   const unsigned int n_dims = _training_params.sizes()[1];
      74             : 
      75             :   mooseAssert(x.size() == n_dims,
      76             :               "Number of parameters provided for evaluation does not match number of parameters "
      77             :               "used for training.");
      78      134001 :   const unsigned int n_outputs = _gp.getCovarFunction().numOutputs();
      79             : 
      80      268002 :   y = std::vector<Real>(n_outputs, 0.0);
      81      268002 :   std = std::vector<Real>(n_outputs, 0.0);
      82             : 
      83      134001 :   const auto options = _training_params.options().dtype(at::kDouble);
      84             : 
      85      134001 :   torch::Tensor test_points = torch::empty({1, n_dims}, at::kDouble);
      86      134001 :   auto points_accessor = test_points.accessor<Real, 2>();
      87      402256 :   for (unsigned int ii = 0; ii < n_dims; ++ii)
      88      268255 :     points_accessor[0][ii] = x[ii];
      89      134001 :   test_points = test_points.to(options.device());
      90             : 
      91      134001 :   _gp.getParamStandardizer().getStandardized(test_points);
      92             : 
      93             :   torch::Tensor K_train_test =
      94      134001 :       torch::empty({_training_params.sizes()[0] * n_outputs, n_outputs}, options);
      95             : 
      96      134001 :   _gp.getCovarFunction().computeCovarianceMatrix(
      97             :       K_train_test, _training_params, test_points, false);
      98      134001 :   torch::Tensor K_test = torch::empty({n_outputs, n_outputs}, options);
      99      134001 :   _gp.getCovarFunction().computeCovarianceMatrix(K_test, test_points, test_points, true);
     100             : 
     101             :   //  Compute the predicted mean value (centered)
     102             :   torch::Tensor pred_value = torch::transpose(
     103      134001 :       torch::mm(torch::transpose(K_train_test, 0, 1), _gp.getKResultsSolve()), 0, 1);
     104             : 
     105             :   //   De-center/scale the value and store for return
     106      134001 :   _gp.getDataStandardizer().getDestandardized(pred_value);
     107             : 
     108             :   torch::Tensor pred_var =
     109      134001 :       K_test - torch::mm(torch::transpose(K_train_test, 0, 1),
     110      402003 :                          torch::cholesky_solve(K_train_test, _gp.getKCholeskyDecomp()));
     111             : 
     112             :   // Only the marginal variances are returned. Clamp tiny negative roundoff before sqrt.
     113             :   torch::Tensor std_dev_vec =
     114      134001 :       torch::sqrt(torch::clamp_min(torch::diagonal(pred_var), 0.0)).unsqueeze(0);
     115      134001 :   _gp.getDataStandardizer().getDescaled(std_dev_vec);
     116      134001 :   const auto std_dev_cpu = LibtorchUtils::toCPUContiguous(std_dev_vec);
     117      134001 :   const auto pred_value_cpu = LibtorchUtils::toCPUContiguous(pred_value);
     118      134001 :   auto std_accessor = std_dev_cpu.accessor<Real, 2>();
     119      134001 :   auto pred_value_accessor = pred_value_cpu.accessor<Real, 2>();
     120             : 
     121      268077 :   for (const auto output_i : make_range(n_outputs))
     122             :   {
     123      134076 :     y[output_i] = pred_value_accessor[0][output_i];
     124      134076 :     std[output_i] = std_accessor[0][output_i];
     125             :   }
     126      134001 : }
     127             : 
     128             : #endif