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 "LoadCovarianceDataAction.h"
      13             : #include "CovarianceFunctionBase.h"
      14             : #include "GaussianProcessSurrogate.h"
      15             : #include "FEProblem.h"
      16             : #include "StochasticToolsApp.h"
      17             : 
      18             : registerMooseAction("StochasticToolsApp", LoadCovarianceDataAction, "load_covariance_data");
      19             : 
      20             : namespace
      21             : {
      22             : 
      23             : using HyperParameterMap = CovarianceFunctionBase::HyperParameterMap;
      24             : 
      25             : std::vector<Real>
      26          16 : exportHyperParameter(const torch::Tensor & tensor)
      27             : {
      28          16 :   const auto flattened = tensor.reshape({-1}).contiguous();
      29          48 :   return {flattened.data_ptr<Real>(), flattened.data_ptr<Real>() + flattened.numel()};
      30             : }
      31             : 
      32             : void
      33          32 : assignRequiredHyperParameter(InputParameters & params,
      34             :                              const std::string & param_name,
      35             :                              const torch::Tensor & tensor)
      36             : {
      37          32 :   if (params.have_parameter<Real>(param_name))
      38             :   {
      39          16 :     if (!CovarianceFunctionBase::isScalarHyperParameter(tensor))
      40           0 :       mooseError("Expected scalar hyperparameter for ", param_name, ".");
      41          16 :     params.set<Real>(param_name) = tensor.item<Real>();
      42             :   }
      43          16 :   else if (params.have_parameter<unsigned int>(param_name))
      44             :   {
      45           0 :     if (!CovarianceFunctionBase::isScalarHyperParameter(tensor))
      46           0 :       mooseError("Expected scalar hyperparameter for ", param_name, ".");
      47           0 :     params.set<unsigned int>(param_name) = cast_int<unsigned int>(tensor.item<Real>());
      48             :   }
      49          16 :   else if (params.have_parameter<std::vector<Real>>(param_name))
      50             :   {
      51          16 :     if (!CovarianceFunctionBase::isVectorHyperParameter(tensor))
      52           0 :       mooseError("Expected vector hyperparameter for ", param_name, ".");
      53          32 :     params.set<std::vector<Real>>(param_name) = exportHyperParameter(tensor);
      54             :   }
      55             :   else
      56           0 :     mooseError("Unsupported hyperparameter type for ", param_name, ".");
      57          32 : }
      58             : 
      59             : void
      60          48 : loadRequiredHyperParameter(InputParameters & params,
      61             :                            const UserObjectName & object_name,
      62             :                            const std::string & param_name,
      63             :                            const HyperParameterMap & hyperparameters)
      64             : {
      65         144 :   const auto expected_name = std::string(object_name) + ":" + param_name;
      66             :   const auto hyperparam_it = hyperparameters.find(expected_name);
      67          48 :   if (hyperparam_it != hyperparameters.end())
      68          32 :     assignRequiredHyperParameter(params, param_name, hyperparam_it->second);
      69          48 : }
      70             : 
      71             : } // namespace
      72             : 
      73             : InputParameters
      74        3008 : LoadCovarianceDataAction::validParams()
      75             : {
      76        3008 :   InputParameters params = Action::validParams();
      77        3008 :   params.addClassDescription("Calls load method on SurrogateModel objects contained within the "
      78             :                              "`[Surrogates]` input block, if a filename is given.");
      79        3008 :   return params;
      80           0 : }
      81             : 
      82        3008 : LoadCovarianceDataAction::LoadCovarianceDataAction(const InputParameters & params) : Action(params)
      83             : {
      84        3008 : }
      85             : 
      86             : void
      87        3006 : LoadCovarianceDataAction::act()
      88             : {
      89             :   std::vector<SurrogateModel *> objects;
      90        6012 :   _app.theWarehouse().query().condition<AttribSystem>("SurrogateModel").queryInto(objects);
      91        3302 :   for (auto model_ptr : objects)
      92             :   {
      93         296 :     auto * gp_gen = dynamic_cast<GaussianProcessSurrogate *>(model_ptr);
      94         760 :     if (gp_gen && model_ptr->isParamValid("filename"))
      95          16 :       load(*gp_gen);
      96             :   }
      97        3006 : }
      98             : 
      99             : void
     100          16 : LoadCovarianceDataAction::load(GaussianProcessSurrogate & model)
     101             : {
     102             :   // We grab all the necessary information that is needed to reconstruct the
     103             :   // covariance structure for the GP
     104             :   const std::string & covar_type = model.getGP().getCovarType();
     105             :   const std::string & covar_name = model.getGP().getCovarName();
     106             :   const std::map<UserObjectName, std::string> & dep_covar_types =
     107             :       model.getGP().getDependentCovarTypes();
     108             :   const std::vector<UserObjectName> & dep_covar_names = model.getGP().getDependentCovarNames();
     109             : 
     110             :   // This is for the covariance on the very top, the lower-level covariances are
     111             :   // all assumed to have num_outputs=1.
     112          16 :   const unsigned int num_outputs = model.getGP().getCovarNumOutputs();
     113             :   const HyperParameterMap & hyperparameters = model.getGP().getHyperParamMap();
     114             : 
     115             :   // We start by creating and loading the lower-level covariances if they need
     116             :   // to be present. Right now we can only load a complex covariance which has
     117             :   // a one-level dependency depth.
     118             :   // TODO: Extend this to arbitrary dependency depths. Maybe we could use a graph.
     119          24 :   for (const auto & it : dep_covar_types)
     120             :   {
     121           8 :     const auto & name = it.first;
     122           8 :     const auto & type = it.second;
     123           8 :     InputParameters covar_params = _factory.getValidParams(type);
     124             : 
     125             :     // We make sure that every required parameter is added so that the object
     126             :     // can be constructed. The non-required hyperparameters (if present in the
     127             :     // parameter maps) will be inserted later.
     128           8 :     const auto param_list = covar_params.getParametersList();
     129         184 :     for (const auto & param : param_list)
     130         176 :       if (covar_params.isParamRequired(param))
     131          16 :         loadRequiredHyperParameter(covar_params, name, param, hyperparameters);
     132             : 
     133          16 :     _problem->addObject<CovarianceFunctionBase>(type, name, covar_params, /*threaded=*/false);
     134           8 :   }
     135             : 
     136          16 :   InputParameters covar_params = _factory.getValidParams(covar_type);
     137          16 :   covar_params.set<unsigned int>("num_outputs") = num_outputs;
     138          16 :   covar_params.set<std::vector<UserObjectName>>("covariance_functions") = dep_covar_names;
     139             : 
     140          16 :   const auto param_list = covar_params.getParametersList();
     141         352 :   for (const auto & param : param_list)
     142             :     // We make sure that every required parameter is added so that the object
     143             :     // can be constructed. The non-required hyperparameters (if present in the
     144             :     // parameter maps) will be inserted later.
     145         336 :     if (covar_params.isParamRequired(param))
     146          64 :       loadRequiredHyperParameter(covar_params, covar_name, param, hyperparameters);
     147             : 
     148          16 :   auto covar_object = _problem->addObject<CovarianceFunctionBase>(
     149          32 :       covar_type, covar_name, covar_params, /* threaded = */ false);
     150          16 :   covar_object[0]->loadHyperParamMap(hyperparameters);
     151             : 
     152          16 :   model.setupCovariance(covar_name);
     153          32 : }
     154             : 
     155             : #endif