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             : #pragma once
      12             : 
      13             : #include "Standardizer.h"
      14             : 
      15             : #include "CovarianceFunctionBase.h"
      16             : 
      17             : #include "LibtorchUtils.h"
      18             : 
      19             : namespace StochasticTools
      20             : {
      21             : 
      22             : /**
      23             :  * Utility class dedicated to hold structures and functions commont to
      24             :  * Gaussian Processes. It can be used to standardize parameters, manipulate
      25             :  * covariance data and compute additional stored matrices.
      26             :  */
      27             : class GaussianProcess
      28             : {
      29             : public:
      30             :   using HyperParameterMap = CovarianceFunctionBase::HyperParameterMap;
      31             : 
      32             :   GaussianProcess();
      33             : 
      34             :   enum class OptimizerType
      35             :   {
      36             :     Adam,
      37             :     LegacyAdam
      38             :   };
      39             : 
      40             :   /**
      41             :    * Initializes the most important structures in the Gaussian Process: the
      42             :    * covariance function and a tuning map which is used if the user requires
      43             :    * parameter tuning.
      44             :    * @param covariance_function Pointer to the covariance function that
      45             :    *                            needs to be used for the Gaussian Process.
      46             :    * @param params_to_tune List of parameters which need to be tuned.
      47             :    * @param min List of lower bounds for the parameter tuning.
      48             :    * @param max List of upper bounds for parameter tuning.
      49             :    */
      50             :   void initialize(CovarianceFunctionBase * covariance_function,
      51             :                   const std::vector<std::string> & params_to_tune,
      52             :                   const std::vector<Real> & min = std::vector<Real>(),
      53             :                   const std::vector<Real> & max = std::vector<Real>());
      54             : 
      55             :   /// Structure containing the optimization options for
      56             :   /// hyperparameter-tuning
      57             :   struct GPOptimizerOptions
      58             :   {
      59             :     /**
      60             :      * Construct a new GPOptimizerOptions object using
      61             :      * input parameters that will control the optimization
      62             :      * @param show_every_nth_iteration To show the loss value at every n-th iteration, if set to 0,
      63             :      * nothing is displayed
      64             :      * @param num_iter The number of iterations we want in the optimization of the GP
      65             :      * @param batch_size The number of samples in each batch
      66             :      * @param learning_rate The learning rate for parameter updates
      67             :      * @param b1 Tuning constant for the Adam algorithm
      68             :      * @param b2 Tuning constant for the Adam algorithm
      69             :      * @param eps Tuning constant for the Adam algorithm
      70             :      * @param lambda Legacy MOOSE shrink constant for the Adam algorithm
      71             :      * @param optimizer_type The Adam optimizer mode to use
      72             :      */
      73             :     GPOptimizerOptions(const unsigned int show_every_nth_iteration = 0,
      74             :                        const unsigned int num_iter = 1000,
      75             :                        const unsigned int batch_size = 0,
      76             :                        const Real learning_rate = 1e-3,
      77             :                        const Real b1 = 0.9,
      78             :                        const Real b2 = 0.999,
      79             :                        const Real eps = 1e-7,
      80             :                        const Real lambda = 1e-4,
      81             :                        const OptimizerType optimizer_type = OptimizerType::Adam);
      82             : 
      83             :     /// Switch to enable verbose output for parameter tuning at every n-th iteration
      84             :     const unsigned int show_every_nth_iteration = 0;
      85             :     /// The number of iterations for Adam optimizer
      86             :     const unsigned int num_iter = 1000;
      87             :     /// The batch size for Adam optimizer
      88             :     const unsigned int batch_size = 0;
      89             :     /// The learning rate for Adam optimizer
      90             :     const Real learning_rate = 1e-3;
      91             :     /// Tuning parameter from the paper
      92             :     const Real b1 = 0.9;
      93             :     /// Tuning parameter from the paper
      94             :     const Real b2 = 0.999;
      95             :     /// Tuning parameter from the paper
      96             :     const Real eps = 1e-7;
      97             :     /// Legacy MOOSE shrink parameter
      98             :     const Real lambda = 1e-4;
      99             :     /// Adam optimizer mode to use
     100             :     const OptimizerType optimizer_type = OptimizerType::Adam;
     101             :   };
     102             :   /**
     103             :    * Sets up the covariance matrix given data and optimization options.
     104             :    * @param training_params The training parameter values (x values) for the
     105             :    *                        covariance matrix.
     106             :    * @param training_data The training data (y values) for the inversion of the
     107             :    *                      covariance matrix.
     108             :    * @param opts The optimizer options.
     109             :    */
     110             :   void setupCovarianceMatrix(const torch::Tensor & training_params,
     111             :                              const torch::Tensor & training_data,
     112             :                              const GPOptimizerOptions & opts);
     113             : 
     114             :   /**
     115             :    * Sets up the Cholesky decomposition and inverse action of the covariance matrix.
     116             :    * @param input The vector/matrix which right multiples the inverse of the covariance matrix.
     117             :    */
     118             :   void setupStoredMatrices(const torch::Tensor & input);
     119             : 
     120             :   /**
     121             :    * Finds and links the covariance function to this object. Used mainly in the
     122             :    * covariance data action.
     123             :    * @param covariance_function Pointer to the covariance function that
     124             :    *                            needs to be used for the Gaussian Process.
     125             :    */
     126             :   void linkCovarianceFunction(CovarianceFunctionBase * covariance_function);
     127             : 
     128             :   /**
     129             :    * Sets up the tuning map which is used if the user requires parameter tuning.
     130             :    * @param params_to_tune List of parameters which need to be tuned.
     131             :    * @param min List of lower bounds for the parameter tuning.
     132             :    * @param max List of upper bounds for parameter tuning.
     133             :    */
     134             :   void generateTuningMap(const std::vector<std::string> & params_to_tune,
     135             :                          const std::vector<Real> & min = std::vector<Real>(),
     136             :                          const std::vector<Real> & max = std::vector<Real>());
     137             : 
     138             :   /**
     139             :    * Standardizes the vector of input parameters (x values).
     140             :    * @param parameters The vector/matrix of input data.
     141             :    * @param keep_moments If previously computed or new moments are to be used.
     142             :    */
     143             :   void standardizeParameters(torch::Tensor & parameters, bool keep_moments = false);
     144             : 
     145             :   /**
     146             :    * Standardizes the vector of responses (y values).
     147             :    * @param data The vector/matrix of input data.
     148             :    * @param keep_moments If previously computed or new moments are to be used.
     149             :    */
     150             :   void standardizeData(torch::Tensor & data, bool keep_moments = false);
     151             : 
     152             :   // Tune hyperparameters using Adam with manually supplied gradients
     153             :   void tuneHyperParamsAdam(const torch::Tensor & training_params,
     154             :                            const torch::Tensor & training_data,
     155             :                            const GPOptimizerOptions & opts);
     156             : 
     157             :   // Computes the loss function
     158             :   Real getLoss(torch::Tensor & inputs, torch::Tensor & outputs);
     159             : 
     160             :   // Computes Gradient of the loss function
     161             :   std::vector<Real> getGradient(torch::Tensor & inputs) const;
     162             : 
     163             :   /// Function used to convert the hyperparameter map in this object to
     164             :   /// a flat vector
     165             :   void mapToVec(
     166             :       const std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> &
     167             :           tuning_data,
     168             :       const HyperParameterMap & hyperparam_map,
     169             :       std::vector<Real> & vec) const;
     170             : 
     171             :   /// Function used to convert the vector back to the hyperparameter map
     172             :   void vecToMap(
     173             :       const std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> &
     174             :           tuning_data,
     175             :       HyperParameterMap & hyperparam_map,
     176             :       const std::vector<Real> & vec) const;
     177             : 
     178             :   /// @{
     179             :   /**
     180             :    * Get constant reference to the contained structures
     181             :    */
     182      134001 :   const StochasticTools::Standardizer & getParamStandardizer() const { return _param_standardizer; }
     183      268002 :   const StochasticTools::Standardizer & getDataStandardizer() const { return _data_standardizer; }
     184             :   const torch::Tensor & getK() const { return _K; }
     185      134001 :   const torch::Tensor & getKResultsSolve() const { return _K_results_solve; }
     186      134001 :   const torch::Tensor & getKCholeskyDecomp() const { return _K_cho_decomp; }
     187      402131 :   const CovarianceFunctionBase & getCovarFunction() const { return *_covariance_function; }
     188          16 :   const CovarianceFunctionBase * getCovarFunctionPtr() const { return _covariance_function; }
     189          16 :   const std::string & getCovarType() const { return _covar_type; }
     190          16 :   const std::string & getCovarName() const { return _covar_name; }
     191             :   const std::vector<UserObjectName> & getDependentCovarNames() const
     192             :   {
     193          16 :     return _dependent_covar_names;
     194             :   }
     195             :   const std::map<UserObjectName, std::string> & getDependentCovarTypes() const
     196             :   {
     197             :     return _dependent_covar_types;
     198             :   }
     199             :   const unsigned int & getCovarNumOutputs() const { return _num_outputs; }
     200             :   const unsigned int & getNumTunableParams() const { return _num_tunable; }
     201          16 :   const HyperParameterMap & getHyperParamMap() const { return _hyperparam_map; }
     202          86 :   const std::vector<Real> & getLengthScales() const { return _length_scales; }
     203             :   ///@}
     204             : 
     205             :   /// @{
     206             :   /**
     207             :    * Get non-constant reference to the contained structures (if they need to be modified from the
     208             :    * utside)
     209             :    */
     210          39 :   StochasticTools::Standardizer & paramStandardizer() { return _param_standardizer; }
     211          39 :   StochasticTools::Standardizer & dataStandardizer() { return _data_standardizer; }
     212          39 :   torch::Tensor & K() { return _K; }
     213          39 :   torch::Tensor & KResultsSolve() { return _K_results_solve; }
     214          39 :   torch::Tensor & KCholeskyDecomp() { return _K_cho_decomp; }
     215             :   CovarianceFunctionBase * covarFunctionPtr() { return _covariance_function; }
     216             :   CovarianceFunctionBase & covarFunction() { return *_covariance_function; }
     217          39 :   std::string & covarType() { return _covar_type; }
     218          39 :   std::string & covarName() { return _covar_name; }
     219          39 :   std::map<UserObjectName, std::string> & dependentCovarTypes() { return _dependent_covar_types; }
     220          39 :   std::vector<UserObjectName> & dependentCovarNames() { return _dependent_covar_names; }
     221          39 :   unsigned int & covarNumOutputs() { return _num_outputs; }
     222             :   std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> & tuningData()
     223             :   {
     224             :     return _tuning_data;
     225             :   }
     226          39 :   HyperParameterMap & hyperparamMap() { return _hyperparam_map; }
     227             :   std::vector<Real> & lengthScales() { return _length_scales; }
     228             :   ///@}
     229             : 
     230             : protected:
     231             :   /// Covariance function object
     232             :   CovarianceFunctionBase * _covariance_function = nullptr;
     233             : 
     234             :   /// Contains tuning inforation. Index of hyperparam, size, and min/max bounds
     235             :   std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> _tuning_data;
     236             : 
     237             :   /// Number of tunable hyperparameters
     238             :   unsigned int _num_tunable = 0;
     239             : 
     240             :   /// Type of covariance function used for this GP
     241             :   std::string _covar_type;
     242             : 
     243             :   /// The name of the covariance function used in this GP
     244             :   std::string _covar_name;
     245             : 
     246             :   /// The names of the covariance functions the used covariance function depends on
     247             :   std::vector<UserObjectName> _dependent_covar_names;
     248             : 
     249             :   /// The types of the covariance functions the used covariance function depends on
     250             :   std::map<UserObjectName, std::string> _dependent_covar_types;
     251             : 
     252             :   /// The number of outputs of the GP
     253             :   unsigned int _num_outputs = 0;
     254             : 
     255             :   /// Hyperparameters. Stored as tensors for use in surrogate reload/reporting.
     256             :   HyperParameterMap _hyperparam_map;
     257             : 
     258             :   /// Standardizer for use with params (x)
     259             :   StochasticTools::Standardizer _param_standardizer;
     260             : 
     261             :   /// Standardizer for use with data (y)
     262             :   StochasticTools::Standardizer _data_standardizer;
     263             : 
     264             :   /// An _n_sample by _n_sample covariance matrix constructed from the selected kernel function
     265             :   torch::Tensor _K;
     266             : 
     267             :   /// A solve of Ax=b via Cholesky.
     268             :   torch::Tensor _K_results_solve;
     269             : 
     270             :   /// Cholesky decomposition libtorch tensor object
     271             :   torch::Tensor _K_cho_decomp;
     272             : 
     273             :   /// The batch size for Adam optimization
     274             :   unsigned int _batch_size = 0;
     275             : 
     276             :   /// To return the GP length scales for active learning
     277             :   std::vector<Real> _length_scales;
     278             : };
     279             : 
     280             : } // StochasticTools namespac
     281             : 
     282             : template <>
     283             : void dataStore(std::ostream & stream, StochasticTools::GaussianProcess & gp_utils, void * context);
     284             : template <>
     285             : void dataLoad(std::istream & stream, StochasticTools::GaussianProcess & gp_utils, void * context);
     286             : 
     287             : #endif