NEML2Utils.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 <string>
#include <vector>
#include <utility>
#include <algorithm>
#include <cstring>
#include <map>
 
#ifdef NEML2_ENABLED
 
// NEML2 v3's C++ API exchanges plain at::Tensor (libtorch/ATen); the rich neml2::Tensor /
// typed-tensor library is no longer shipped, so we work with ATen directly here.
#include <ATen/ATen.h>
#include <ATen/Parallel.h> // at::get_num_threads / at::set_num_threads
#include "RankTwoTensor.h"
#include "RankFourTensor.h"
#include "SymmetricRankTwoTensor.h"
#include "SymmetricRankFourTensor.h"
#include "MaterialProperty.h"
 
#endif
 
class MooseObject;
class Action;
class ConsoleStream;
 
namespace NEML2Utils
{
#ifdef NEML2_ENABLED
 
enum class MOOSEIOType
{
  // unbatched
  TIME,
  SCALAR,
  // batched
  FUNCTION,
  VARIABLE,
  MATERIAL
};
 
std::string stringify(MOOSEIOType type);
 
inline std::pair<std::string, int>
parseLag(const std::string & name)
{
  const auto pos = name.rfind('~');
  if (pos == std::string::npos)
    return {name, 0};
  const auto suffix = name.substr(pos + 1);
  if (suffix.empty() || suffix.find_first_not_of("0123456789") != std::string::npos)
    return {name, 0};
  return {name.substr(0, pos), std::stoi(suffix)};
}
 
inline std::string
lagName(const std::string & base, int lag)
{
  return lag == 0 ? base : base + "~" + std::to_string(lag);
}
 
inline bool
contains(const std::vector<std::string> & v, const std::string & s)
{
  return std::find(v.begin(), v.end(), s) != v.end();
}
 
template <typename T>
struct Layout
{
};
template <>
struct Layout<Real>
{
  static constexpr std::array<int64_t, 0> shape{};
  static constexpr std::array<int64_t, 1> strides{1};
};
template <>
struct Layout<RealVectorValue>
{
  static constexpr std::array<int64_t, 1> shape{3};
  static constexpr std::array<int64_t, 2> strides{3, 1};
};
template <>
struct Layout<RankTwoTensor>
{
  static constexpr std::array<int64_t, 2> shape{3, 3};
  static constexpr std::array<int64_t, 3> strides{9, 3, 1};
};
template <>
struct Layout<SymmetricRankTwoTensor>
{
  static constexpr std::array<int64_t, 1> shape{6};
  static constexpr std::array<int64_t, 2> strides{6, 1};
};
template <>
struct Layout<RankFourTensor>
{
  static constexpr std::array<int64_t, 4> shape{3, 3, 3, 3};
  static constexpr std::array<int64_t, 5> strides{81, 27, 9, 3, 1};
};
template <>
struct Layout<SymmetricRankFourTensor>
{
  static constexpr std::array<int64_t, 2> shape{6, 6};
  static constexpr std::array<int64_t, 3> strides{36, 6, 1};
};
 
template <typename T>
at::Tensor
fromBlob(const std::vector<T> & data)
{
  // Full shape: leading batch dim followed by the base shape.
  std::vector<int64_t> shape;
  shape.reserve(1 + Layout<T>::shape.size());
  shape.push_back(static_cast<int64_t>(data.size()));
  shape.insert(shape.end(), Layout<T>::shape.begin(), Layout<T>::shape.end());
  // The const_cast is fine because torch works with non-const ptr so that it can optionally handle
  // deallocation. But we are not going to let torch do that.
  return at::from_blob(const_cast<T *>(data.data()), shape, at::TensorOptions().dtype(at::kDouble));
}
 
template <typename T>
void
copyTensorToMOOSEData(const at::Tensor & src, T & dest)
{
  if (src.dtype() != at::kDouble)
    mooseError(
        "Cannot copy at::Tensor with dtype ", src.dtype(), " into ", demangle(typeid(T).name()));
  if (src.numel() != Layout<T>::strides[0])
    mooseError("Cannot copy at::Tensor with shape ",
               src.sizes(),
               " into ",
               demangle(typeid(T).name()),
               " with different number of elements.");
  auto dest_tensor = at::from_blob(
      reinterpret_cast<Real *>(&dest), Layout<T>::shape, at::TensorOptions().dtype(at::kDouble));
  dest_tensor.copy_(src.reshape(Layout<T>::shape));
}
 
template <typename T>
void
copyBlobToMOOSEData(const Real * src, T & dest)
{
  std::memcpy(reinterpret_cast<Real *>(&dest), src, Layout<T>::strides[0] * sizeof(Real));
}
 
class ScopedNumThreads
{
public:
  explicit ScopedNumThreads(int n) : _saved(at::get_num_threads()), _changed(n > 0 && n != _saved)
  {
    if (_changed)
      at::set_num_threads(n);
  }
  ~ScopedNumThreads()
  {
    if (_changed)
      at::set_num_threads(_saved);
  }
  ScopedNumThreads(const ScopedNumThreads &) = delete;
  ScopedNumThreads & operator=(const ScopedNumThreads &) = delete;
 
private:
  const int _saved;
  const bool _changed;
};
 
class ScopedDefaultDtype
{
public:
  ScopedDefaultDtype() : _saved(at::get_default_dtype())
  {
    at::set_default_dtype(caffe2::TypeMeta::Make<double>());
  }
  ~ScopedDefaultDtype() { at::set_default_dtype(_saved); }
  ScopedDefaultDtype(const ScopedDefaultDtype &) = delete;
  ScopedDefaultDtype & operator=(const ScopedDefaultDtype &) = delete;
 
private:
  const caffe2::TypeMeta _saved;
};
 
static std::string NEML2_help_message = R""""(
==============================================================================
To debug NEML2 related issues:
1. Build and run MOOSE in dbg mode.
2. Re-run the simulation using the dbg executable, and often times
   NEML2 will provide a more helpful error message.
3. If the error message is not helpful, or if there is still no error message,
   run the simulation through a debugger: See
   https://mooseframework.inl.gov/application_development/debugging.html
4. If the issue is due to a NEML2 bug, feel free to report it at
   https://github.com/applied-material-modeling/neml2/issues
==============================================================================
)"""";
 
void redirectLogsToConsole(const ConsoleStream & console);
 
void dumpInputsToTorchScript(const std::map<std::string, at::Tensor> & inputs,
                             const std::string & filename);
 
#endif // NEML2_ENABLED
 
std::string docstring(const std::string & desc);
 
void assertNEML2Enabled();
 
} // namespace NEML2Utils