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 LIBTORCH_ENABLED
      11             : 
      12             : #include "TorchScriptMaterial.h"
      13             : 
      14             : registerMooseObject("MooseApp", TorchScriptMaterial);
      15             : 
      16             : InputParameters
      17        1571 : TorchScriptMaterial::validParams()
      18             : {
      19        1571 :   InputParameters params = Material::validParams();
      20        1571 :   params.addClassDescription(
      21             :       "Material object which relies on the evaluation of a TorchScript module.");
      22        1571 :   params.addRequiredParam<std::vector<std::string>>(
      23             :       "prop_names", "The names of the properties this material will generate.");
      24        1571 :   params.addRequiredParam<std::vector<PostprocessorName>>(
      25             :       "input_names", "The input parameters for the neural network.");
      26        1571 :   params.addRequiredParam<UserObjectName>(
      27             :       "torch_script_userobject",
      28             :       "The name of the user object which contains the torch script module.");
      29             : 
      30        1571 :   return params;
      31           0 : }
      32             : 
      33           8 : TorchScriptMaterial::TorchScriptMaterial(const InputParameters & parameters)
      34             :   : Material(parameters),
      35           8 :     _prop_names(getParam<std::vector<std::string>>("prop_names")),
      36           8 :     _num_props(_prop_names.size()),
      37           8 :     _input_names(getParam<std::vector<PostprocessorName>>("input_names")),
      38           8 :     _num_inputs(_input_names.size()),
      39           8 :     _torch_script_userobject(getUserObject<TorchScriptUserObject>("torch_script_userobject")),
      40           8 :     _input_tensor(torch::zeros(
      41           8 :         {1, _num_inputs},
      42          24 :         torch::TensorOptions().dtype(torch::kFloat64).device(_app.getLibtorchDevice())))
      43             : {
      44           8 :   if (!_num_props)
      45           1 :     paramError("prop_names", "Must declare at least one property!");
      46             : 
      47           7 :   if (!_num_inputs)
      48           1 :     paramError("input_names", "Must declare at least one input to the neural net!");
      49             : 
      50          24 :   for (const auto & input_name : _input_names)
      51          18 :     _module_inputs.push_back(&getPostprocessorValueByName(input_name));
      52             : 
      53          15 :   for (const auto & prop_name : _prop_names)
      54           9 :     _properties.push_back(&declareGenericProperty<Real, false>(prop_name));
      55           6 : }
      56             : 
      57             : void
      58           0 : TorchScriptMaterial::initQpStatefulProperties()
      59             : {
      60           0 :   computeQpValues();
      61           0 : }
      62             : 
      63             : void
      64         221 : TorchScriptMaterial::computeQpProperties()
      65             : {
      66         221 :   computeQpValues();
      67         220 : }
      68             : 
      69             : void
      70         221 : TorchScriptMaterial::computeQpValues()
      71             : {
      72         221 :   auto input_accessor = _input_tensor.accessor<Real, 2>();
      73         884 :   for (unsigned int input_i = 0; input_i < _num_inputs; ++input_i)
      74         663 :     input_accessor[0][input_i] = (*_module_inputs[input_i]);
      75             : 
      76         221 :   const auto output = _torch_script_userobject.evaluate(_input_tensor);
      77         221 :   if (_num_props != output.numel())
      78           2 :     mooseError("The tensor needs to be the same length (right now ",
      79           1 :                output.numel(),
      80             :                ") as the number of properties (right now ",
      81           1 :                _num_props,
      82             :                ")!");
      83             : 
      84         220 :   const auto output_accessor = output.accessor<Real, 2>();
      85         440 :   for (unsigned int prop_i = 0; prop_i < _num_props; ++prop_i)
      86         220 :     (*_properties[prop_i])[_qp] = output_accessor[0][prop_i];
      87         220 : }
      88             : 
      89             : #endif