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             : 
      10             : #ifdef MOOSE_LIBTORCH_ENABLED
      11             : 
      12             : #include "LibtorchANNSurrogate.h"
      13             : 
      14             : registerMooseObject("StochasticToolsApp", LibtorchANNSurrogate);
      15             : 
      16             : InputParameters
      17          80 : LibtorchANNSurrogate::validParams()
      18             : {
      19          80 :   InputParameters params = SurrogateModel::validParams();
      20          80 :   params.addClassDescription("Surrogate that evaluates a feedforward artificial neural net. ");
      21          80 :   return params;
      22           0 : }
      23             : 
      24          40 : LibtorchANNSurrogate::LibtorchANNSurrogate(const InputParameters & parameters)
      25             :   : SurrogateModel(parameters),
      26          40 :     _nn(getModelData<std::shared_ptr<Moose::LibtorchArtificialNeuralNet>>("nn")),
      27          80 :     _input_standardizer(getModelData<StochasticTools::Standardizer>("input_standardizer")),
      28         120 :     _output_standardizer(getModelData<StochasticTools::Standardizer>("output_standardizer"))
      29             : {
      30             :   // We check if MOOSE is compiled with torch, if not this throws an error
      31          40 :   StochasticToolsApp::requiresTorch(*this);
      32          40 : }
      33             : 
      34             : Real
      35         310 : LibtorchANNSurrogate::evaluate(const std::vector<Real> & x) const
      36             : {
      37             :   Real val(0.0);
      38             : 
      39             :   // Check whether input point has same dimensionality as training data
      40             :   mooseAssert(_nn->numInputs() == x.size(),
      41             :               "Input point does not match dimensionality of training data.");
      42             : 
      43         310 :   std::vector<Real> converted_input(x.size(), 0);
      44         310 :   const auto & input_mean = _input_standardizer.getMean();
      45             :   const auto & input_std = _input_standardizer.getStdDev();
      46             : 
      47         310 :   auto input_mean_accessor = input_mean.accessor<Real, 1>();
      48         310 :   auto input_std_accessor = input_std.accessor<Real, 1>();
      49             : 
      50             :   mooseAssert((unsigned long int)torch::size(input_mean, 0) == converted_input.size() &&
      51             :                   (unsigned long int)torch::size(input_std, 0) == converted_input.size(),
      52             :               "The input standardizer's dimensions should be the same as the input dimension!");
      53             : 
      54        1390 :   for (auto input_i : index_range(converted_input))
      55        1080 :     converted_input[input_i] =
      56        1080 :         (x[input_i] - input_mean_accessor[input_i]) / input_std_accessor[input_i];
      57             : 
      58             :   torch::Tensor x_tf =
      59         620 :       torch::tensor(torch::ArrayRef<Real>(converted_input.data(), converted_input.size()))
      60         310 :           .to(at::kDouble);
      61             : 
      62         310 :   const auto & output_mean = _output_standardizer.getMean();
      63             :   const auto & output_std = _output_standardizer.getStdDev();
      64             : 
      65         310 :   auto output_mean_accessor = output_mean.accessor<Real, 1>();
      66         310 :   auto output_std_accessor = output_std.accessor<Real, 1>();
      67             : 
      68             :   mooseAssert(output_mean.sizes()[0] == 1 && output_std.sizes()[0] == 1,
      69             :               "The output standardizer's dimensions should be 1!");
      70             : 
      71             :   // Compute prediction
      72         620 :   val = _nn->forward(x_tf).item<double>();
      73         310 :   val = val * output_std_accessor[0] + output_mean_accessor[0];
      74             : 
      75         310 :   return val;
      76         310 : }
      77             : 
      78             : #endif