NEML2ModelExecutor.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 "NEML2ModelInterface.h"
#include "GeneralUserObject.h"
#include "NEML2BatchIndexGenerator.h"
 
class MOOSEToNEML2;
 
class NEML2ModelExecutor : public NEML2ModelInterface<GeneralUserObject>
{
public:
  static InputParameters actionParams();
 
  static InputParameters validParams();
 
  NEML2ModelExecutor(const InputParameters & params);
 
#ifndef NEML2_ENABLED
  void initialize() override {}
  void execute() override {}
  void finalize() override {}
#else
  void initialize() override;
  void meshChanged() override;
  void execute() override;
  void finalize() override;
 
  void initialSetup() override;
 
  std::size_t getBatchIndex(dof_id_type elem_id) const;
 
  const at::Tensor & getOutput(const std::string & output_name) const;
 
  const at::Tensor & getOutputDerivative(const std::string & output_name,
                                         const std::string & input_name) const;
 
  const at::Tensor & getOutputParameterDerivative(const std::string & output_name,
                                                  const std::string & parameter_name) const;
 
  bool outputReady() const { return _output_ready; }
 
protected:
  virtual void addGatheredVariable(const UserObjectName &, const std::string &);
 
  virtual void addGatheredParameter(const UserObjectName &, const std::string &);
 
  virtual void checkExecutionStage() const final;
 
  virtual void fillInputs();
 
  virtual bool solve(bool compute_derivative);
 
  virtual void extractOutputs(bool compute_derivative);
 
  std::string failedInputDumpName(unsigned int count, const std::string & rank_token) const;
 
  void deviceSynchronize();
 
  void advanceState();
 
  const NEML2BatchIndexGenerator & _batch_index_generator;
 
  const bool _manage_state_advance;
 
  const bool _dump_inputs_on_failure;
 
  bool _output_ready;
 
  bool _derivative_valid;
 
  std::map<std::string, at::Tensor> _in;
 
  std::map<std::string, at::Tensor> _model_params;
 
  std::map<std::string, at::Tensor> _out;
 
  std::map<std::string, at::Tensor> _state_vars;
 
  std::map<std::string, at::Tensor> _zero_seed_vars;
 
  std::map<std::string, std::map<std::string, at::Tensor>> _dout_din;
 
  std::map<std::string, std::map<std::string, at::Tensor>> _dout_dparam;
 
  // set of gathered NEML2 input variables
  std::set<std::string> _gathered_variable_names;
 
  // set of gathered NEML2 model parameters
  std::set<std::string> _gathered_parameter_names;
 
  std::vector<const MOOSEToNEML2 *> _gatherers;
 
  std::vector<const MOOSEToNEML2 *> _param_gatherers;
 
  mutable std::map<std::string, at::Tensor> _retrieved_outputs;
 
  mutable std::map<std::string, std::map<std::string, at::Tensor>> _retrieved_derivatives;
 
  mutable std::map<std::string, std::map<std::string, at::Tensor>> _retrieved_parameter_derivatives;
 
private:
  bool _error;
  std::string _error_message;
  unsigned int _num_failed_dumps;
#endif
};