SurrogateTrainer.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

#pragma once

#include "StochasticToolsApp.h"
#include "GeneralUserObject.h"
#include "LoadSurrogateDataAction.h"
#include "RestartableModelInterface.h"

#include "Sampler.h"
#include "StochasticToolsApp.h"
#include "SurrogateModelInterface.h"
#include "MooseRandom.h"

class TrainingDataBase;
template <typename T>
class TrainingData;

class SurrogateTrainerBase : public GeneralUserObject, public RestartableModelInterface
{
public:
  static InputParameters validParams();
  SurrogateTrainerBase(const InputParameters & parameters);

  virtual void initialize() {}                         // not required, but available
  virtual void finalize() {}                           // not required, but available
  virtual void threadJoin(const UserObject &) final {} // GeneralUserObjects are not threaded
};

class SurrogateTrainer : public SurrogateTrainerBase, public SurrogateModelInterface
{
public:
  static InputParameters validParams();
  SurrogateTrainer(const InputParameters & parameters);

  virtual void initialize() final;
  virtual void execute() final;
  virtual void finalize() final {}

protected:
  /*
   * Setup function called before sampler loop
   */
  virtual void preTrain() {}

  /*
   * Function needed to be overried, called during sampler loop
   */
  virtual void train() {}

  /*
   * Function called after sampler loop, used for mpi communication mainly
   */
  virtual void postTrain() {}

  // TRAINING_DATA_BEGIN

  /*
   * Get a reference to training data given a reporter name
   */
  template <typename T>
  const T & getTrainingData(const ReporterName & rname);

  /*
   * Get a reference to the sampler row data
   */
  const std::vector<Real> & getSamplerData() const { return _row_data; };

  /*
   * Get a reference to the predictor row data
   */
  const std::vector<Real> & getPredictorData() const { return _predictor_data; };

  /*
   * Get current sample size (this is recalculated to reflect the number of skipped samples)
   */
  unsigned int getCurrentSampleSize() const { return _current_sample_size; };

  /*
   * Get current local sample size (recalculated to reflect number of skipped samples)
   */
  unsigned int getLocalSampleSize() const { return _local_sample_size; };

  // TRAINING_DATA_END

  /*
   * Evaluate CV error using _cv_surrogate and appropriate predictor row.
   */
  virtual std::vector<Real> evaluateModelError(const SurrogateModel & surr);

  // TRAINING_DATA_MEMBERS
  Sampler & _sampler;
  dof_id_type _row;
  dof_id_type _local_row;
  const Real * _rval;
  const std::vector<Real> * _rvecval;
  std::vector<const Real *> _pvals;
  std::vector<unsigned int> _pcols;
  unsigned int _n_dims;
  unsigned int & _n_outputs;
  // TRAINING_DATA_MEMBERS_END

private:
  /*
   * Called at the beginning of execute() to make sure values are set properly
   */
  void checkIntegrity() const;

  /*
   * Main model training method - called during crossValidate() and for final model training.
   */
  void executeTraining();

  /*
   * Call if cross-validation is turned on.
   */
  std::vector<Real> crossValidate();

  /*
   * Update predictor row (uses both Sampler and Reporter values, according to _pvals and _pcols)
   */
  void updatePredictorRow();

  std::vector<Real> _row_data;

  std::vector<Real> _predictor_data;

  const bool _skip_unconverged;

  const bool * _converged;

  unsigned int _current_sample_size;

  unsigned int _local_sample_size;

  std::unordered_map<ReporterName, std::shared_ptr<TrainingDataBase>> _training_data;

  /*
   * Variables related to cross validation.
   */
  std::vector<dof_id_type> _skip_indices;
  const MooseEnum & _cv_type;
  const unsigned int & _n_splits;
  const unsigned int & _cv_n_trials;
  const unsigned int & _cv_seed;
  MooseRandom _cv_generator;
  const SurrogateModel * _cv_surrogate;
  const bool _doing_cv;
  std::vector<std::vector<Real>> & _cv_trial_scores;
};

template <typename T>
const T &
SurrogateTrainer::getTrainingData(const ReporterName & rname)
{
  auto it = _training_data.find(rname);
  if (it != _training_data.end())
  {
    auto data = std::dynamic_pointer_cast<TrainingData<T>>(it->second);
    if (!data)
      mooseError("Reporter value ", rname, " already exists but is of different type.");
    return data->get();
  }
  else
  {
    const std::vector<T> & rval = getReporterValueByName<std::vector<T>>(rname);
    _training_data[rname] = std::make_shared<TrainingData<T>>(rval);
    return std::dynamic_pointer_cast<TrainingData<T>>(_training_data[rname])->get();
  }
}

class TrainingDataBase
{
public:
  TrainingDataBase() : _is_distributed(false) {}

  virtual ~TrainingDataBase() = default;

  virtual dof_id_type size() const = 0;
  virtual void setCurrentIndex(dof_id_type index) = 0;
  bool & isDistributed() { return _is_distributed; }

protected:
  bool _is_distributed;
};

template <typename T>
class TrainingData : public TrainingDataBase
{
public:
  TrainingData(const std::vector<T> & vector) : _vector(vector) {}

  virtual dof_id_type size() const override { return _vector.size(); }
  virtual void setCurrentIndex(dof_id_type index) override { _value = _vector[index]; }

  const T & get() const { return _value; }

private:
  const std::vector<T> & _vector;
  T _value;
};