GaussianProcess.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html
#ifdef MOOSE_LIBTORCH_ENABLED

#pragma once

#include "Standardizer.h"

#include "CovarianceFunctionBase.h"

#include "LibtorchUtils.h"

namespace StochasticTools
{

class GaussianProcess
{
public:
  using HyperParameterMap = CovarianceFunctionBase::HyperParameterMap;

  GaussianProcess();

  enum class OptimizerType
  {
    Adam,
    LegacyAdam
  };

  void initialize(CovarianceFunctionBase * covariance_function,
                  const std::vector<std::string> & params_to_tune,
                  const std::vector<Real> & min = std::vector<Real>(),
                  const std::vector<Real> & max = std::vector<Real>());

  struct GPOptimizerOptions
  {
    GPOptimizerOptions(const unsigned int show_every_nth_iteration = 0,
                       const unsigned int num_iter = 1000,
                       const unsigned int batch_size = 0,
                       const Real learning_rate = 1e-3,
                       const Real b1 = 0.9,
                       const Real b2 = 0.999,
                       const Real eps = 1e-7,
                       const Real lambda = 1e-4,
                       const OptimizerType optimizer_type = OptimizerType::Adam);

    const unsigned int show_every_nth_iteration = 0;
    const unsigned int num_iter = 1000;
    const unsigned int batch_size = 0;
    const Real learning_rate = 1e-3;
    const Real b1 = 0.9;
    const Real b2 = 0.999;
    const Real eps = 1e-7;
    const Real lambda = 1e-4;
    const OptimizerType optimizer_type = OptimizerType::Adam;
  };
  void setupCovarianceMatrix(const torch::Tensor & training_params,
                             const torch::Tensor & training_data,
                             const GPOptimizerOptions & opts);

  void setupStoredMatrices(const torch::Tensor & input);

  void linkCovarianceFunction(CovarianceFunctionBase * covariance_function);

  void generateTuningMap(const std::vector<std::string> & params_to_tune,
                         const std::vector<Real> & min = std::vector<Real>(),
                         const std::vector<Real> & max = std::vector<Real>());

  void standardizeParameters(torch::Tensor & parameters, bool keep_moments = false);

  void standardizeData(torch::Tensor & data, bool keep_moments = false);

  // Tune hyperparameters using Adam with manually supplied gradients
  void tuneHyperParamsAdam(const torch::Tensor & training_params,
                           const torch::Tensor & training_data,
                           const GPOptimizerOptions & opts);

  // Computes the loss function
  Real getLoss(torch::Tensor & inputs, torch::Tensor & outputs);

  // Computes Gradient of the loss function
  std::vector<Real> getGradient(torch::Tensor & inputs) const;

  void mapToVec(
      const std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> &
          tuning_data,
      const HyperParameterMap & hyperparam_map,
      std::vector<Real> & vec) const;

  void vecToMap(
      const std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> &
          tuning_data,
      HyperParameterMap & hyperparam_map,
      const std::vector<Real> & vec) const;


  const StochasticTools::Standardizer & getParamStandardizer() const { return _param_standardizer; }
  const StochasticTools::Standardizer & getDataStandardizer() const { return _data_standardizer; }
  const torch::Tensor & getK() const { return _K; }
  const torch::Tensor & getKResultsSolve() const { return _K_results_solve; }
  const torch::Tensor & getKCholeskyDecomp() const { return _K_cho_decomp; }
  const CovarianceFunctionBase & getCovarFunction() const { return *_covariance_function; }
  const CovarianceFunctionBase * getCovarFunctionPtr() const { return _covariance_function; }
  const std::string & getCovarType() const { return _covar_type; }
  const std::string & getCovarName() const { return _covar_name; }
  const std::vector<UserObjectName> & getDependentCovarNames() const
  {
    return _dependent_covar_names;
  }
  const std::map<UserObjectName, std::string> & getDependentCovarTypes() const
  {
    return _dependent_covar_types;
  }
  const unsigned int & getCovarNumOutputs() const { return _num_outputs; }
  const unsigned int & getNumTunableParams() const { return _num_tunable; }
  const HyperParameterMap & getHyperParamMap() const { return _hyperparam_map; }
  const std::vector<Real> & getLengthScales() const { return _length_scales; }


  StochasticTools::Standardizer & paramStandardizer() { return _param_standardizer; }
  StochasticTools::Standardizer & dataStandardizer() { return _data_standardizer; }
  torch::Tensor & K() { return _K; }
  torch::Tensor & KResultsSolve() { return _K_results_solve; }
  torch::Tensor & KCholeskyDecomp() { return _K_cho_decomp; }
  CovarianceFunctionBase * covarFunctionPtr() { return _covariance_function; }
  CovarianceFunctionBase & covarFunction() { return *_covariance_function; }
  std::string & covarType() { return _covar_type; }
  std::string & covarName() { return _covar_name; }
  std::map<UserObjectName, std::string> & dependentCovarTypes() { return _dependent_covar_types; }
  std::vector<UserObjectName> & dependentCovarNames() { return _dependent_covar_names; }
  unsigned int & covarNumOutputs() { return _num_outputs; }
  std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> & tuningData()
  {
    return _tuning_data;
  }
  HyperParameterMap & hyperparamMap() { return _hyperparam_map; }
  std::vector<Real> & lengthScales() { return _length_scales; }

protected:
  CovarianceFunctionBase * _covariance_function = nullptr;

  std::unordered_map<std::string, std::tuple<unsigned int, unsigned int, Real, Real>> _tuning_data;

  unsigned int _num_tunable = 0;

  std::string _covar_type;

  std::string _covar_name;

  std::vector<UserObjectName> _dependent_covar_names;

  std::map<UserObjectName, std::string> _dependent_covar_types;

  unsigned int _num_outputs = 0;

  HyperParameterMap _hyperparam_map;

  StochasticTools::Standardizer _param_standardizer;

  StochasticTools::Standardizer _data_standardizer;

  torch::Tensor _K;

  torch::Tensor _K_results_solve;

  torch::Tensor _K_cho_decomp;

  unsigned int _batch_size = 0;

  std::vector<Real> _length_scales;
};

} // StochasticTools namespac

template <>
void dataStore(std::ostream & stream, StochasticTools::GaussianProcess & gp_utils, void * context);
template <>
void dataLoad(std::istream & stream, StochasticTools::GaussianProcess & gp_utils, void * context);

#endif