TorchFEMUtils.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 <cmath>
#include <ATen/ATen.h>
 
namespace TorchFEM
{
inline at::Tensor
scatter(const at::Tensor & v, const at::Tensor & dof_map)
{
  return v.index({dof_map});
}
 
inline at::Tensor
interpolate(const at::Tensor & elem_dofs, const at::Tensor & basis)
{
  auto ed = elem_dofs;
  while (ed.dim() < basis.dim())
    ed = ed.unsqueeze(-1);
  return (ed * basis).sum(1);
}
 
inline at::Tensor
assemble(const at::Tensor & scattered, const at::Tensor & dof_map, int64_t ndof)
{
  return at::zeros({ndof}, scattered.options())
      .scatter_add_(0, dof_map.flatten(), scattered.flatten());
}
 
inline at::Tensor
fullToMandel(const at::Tensor & full)
{
  const double s2 = std::sqrt(2.0);
  auto c = [&full](int i, int j) { return full.select(-2, i).select(-1, j); };
  return at::stack({c(0, 0), c(1, 1), c(2, 2), s2 * c(1, 2), s2 * c(0, 2), s2 * c(0, 1)}, -1);
}
 
inline at::Tensor
mandelToFull(const at::Tensor & mandel)
{
  const double is2 = 1.0 / std::sqrt(2.0);
  auto m = [&mandel](int k) { return mandel.select(-1, k); };
  auto d0 = m(0), d1 = m(1), d2 = m(2);
  auto o12 = m(3) * is2, o02 = m(4) * is2, o01 = m(5) * is2;
  auto row0 = at::stack({d0, o01, o02}, -1);
  auto row1 = at::stack({o01, d1, o12}, -1);
  auto row2 = at::stack({o02, o12, d2}, -1);
  return at::stack({row0, row1, row2}, -2);
}
} // namespace TorchFEM
 
#endif // MOOSE_LIBTORCH_ENABLED