TorchFEInterpolation.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
 
#ifdef MOOSE_LIBTORCH_ENABLED
 
#include <torch/torch.h>
 
// system includes
#include <vector>
#include <unordered_map>
 
// libmesh includes
#include "MooseTypes.h"
#include "libmesh/fe_type.h"
#include "libmesh/id_types.h"
 
// MOOSE includes
#include "ElementUserObject.h"
#include "TorchAssembly.h"
#include "MooseVariableFE.h"
 
class TorchFEInterpolation : public ElementUserObject
{
public:
  static InputParameters validParams();
 
  TorchFEInterpolation(const InputParameters & parameters);
 
  const at::Tensor & getValue(const std::string & var_name);
 
  const at::Tensor & getGradient(const std::string & var_name);
 
  const at::Tensor & getPhi(const std::string & var_name);
 
  const at::Tensor & getPhiGradient(const std::string & var_name);
 
  const at::Tensor & getDofMap(const std::string & var_name);
 
  const std::vector<dof_id_type> & getGlobalDofMap(const std::string & var_name);
 
  int64_t local_ndof() const;
 
  bool contextUpToDate() const { return _fem_context_up_to_date; }
 
  void invalidateFEMContext();
 
  void invalidateInterpolations();
 
  void initialSetup() override;
  void meshChanged() override;
 
  void initialize() override;
  void execute() override;
  void finalize() override;
  void threadJoin(const UserObject &) override;
 
protected:
  virtual void syncWithMainThread();
  virtual void updateFEMContext();
  virtual void updateDofMap();
  virtual void updatePhi();
  virtual void updateGradPhi();
  virtual void updateInterpolations();
 
  const TorchAssembly & _assembly;
 
  bool _fem_context_up_to_date = false;
 
  bool _interp_up_to_date = false;
 
  std::unordered_map<std::string, at::Tensor> _vars;
 
  std::unordered_map<std::string, at::Tensor> _grad_vars;
 
  std::unordered_map<std::string, const MooseVariableFE<Real> *> _moose_vars;
 
  std::unordered_map<FEType, int64_t> _ndofe;
  std::unordered_map<FEType, const VariablePhiValue *> _phis;
  std::unordered_map<FEType, const VariablePhiGradient *> _grad_phis;
 
  std::unordered_map<std::string, std::vector<int64_t>> _moose_dof_map;
  std::unordered_map<std::string, std::vector<dof_id_type>> _moose_dof_map_global;
  std::unordered_map<FEType, std::vector<Real>> _moose_phi;
  std::unordered_map<FEType, std::vector<Real>> _moose_grad_phi;
 
  std::unordered_map<std::string, at::Tensor> _torch_dof_map;
  std::unordered_map<FEType, at::Tensor> _torch_phi;
  std::unordered_map<FEType, at::Tensor> _torch_grad_phi;
 
  const libMesh::PetscVector<Real> * _petsc_solution;
 
private:
  const MooseVariableFE<Real> * getMOOSEVariable(const std::string & var_name) const;
 
  std::vector<dof_id_type> _dof_indices;
 
  int64_t _local_ndof;
};
 
#endif // MOOSE_LIBTORCH_ENABLED