MeshField_ShapeField.hpp

#ifndef MESHFIELD_SHAPEFIELD_HPP
#define MESHFIELD_SHAPEFIELD_HPP
 
#ifdef MESHFIELDS_ENABLE_CABANA
#include "CabanaController.hpp"
#endif
#include "KokkosController.hpp"
#include "MeshField_Field.hpp"
#include "MeshField_Shape.hpp"
#include <type_traits> //decltype
#include <utility> // std::forward
 
namespace MeshField {
 
struct MeshInfo {
  int numVtx = 0;     // entDim = 0
  int numEdge = 0;    // entDim = 1
  int numTri = 0;     // entDim = 2
  int numQuad = 0;    // entDim = 2
  int numTet = 0;     // entDim = 3
  int numHex = 0;     // entDim = 3
  int numPrism = 0;   // entDim = 3
  int numPyramid = 0; // entDim = 3
  int dim = 0;
 
  int numEntities(Mesh_Topology topo) const {
    switch (topo) {
    case Vertex:
      return numVtx;
    case Edge:
      return numEdge;
    case Triangle:
      return numTri;
    case Quad:
      return numQuad;
    case Tetrahedron:
      return numTet;
    case Hexahedron:
      return numHex;
    case Prism:
      return numPrism;
    case Pyramid:
      return numPyramid;
    }
    return 0;
  }
};
 
template <typename Ctrlr, typename Field> struct FieldWithController {
  Ctrlr ctrlr;
  Field field;
};
 
template <size_t numCompIn, typename Shape, typename... Mixins>
struct ShapeField : public Mixins... {
  Shape shape;
  static const size_t numComp = numCompIn;
  const MeshInfo meshInfo;
  constexpr static auto Order = Shape::Order;
  // Mixins&&... binds only to rvalues/temporaries (braced-init-list call
  // sites like ShapeField(meshInfo, {vtxField}) are fine; passing a named
  // lvalue requires std::move). mixins is moved-from by std::forward below
  // - do not read mixins again after the initializer list.
  ShapeField(const MeshInfo &meshInfoIn, Mixins &&... mixins)
      : meshInfo(meshInfoIn), Mixins( std::forward<Mixins>(mixins) )... {};
};
 
template <typename VtxAccessor, typename EdgeAccessor>
struct QuadraticAccessor {
  constexpr static const Mesh_Topology topo[2] = {Vertex, Edge};
  VtxAccessor vtxField;
  EdgeAccessor edgeField;
  using BaseType = typename VtxAccessor::BaseType;
 
  KOKKOS_FUNCTION
  auto &operator()(int entity, int node, int component, Mesh_Topology t) const {
    if (t != Vertex && t != Edge) {
      Kokkos::printf("%d is not a support topology\n", t);
      assert(false);
    }
    return (t == Vertex) ? vtxField(entity, node, component)
                         : edgeField(entity, node, component);
  }
};
 
template <typename VtxAccessor> struct LinearAccessor {
  constexpr static const Mesh_Topology topo[1] = {Vertex};
  VtxAccessor vtxField;
  using BaseType = typename VtxAccessor::BaseType;
 
  KOKKOS_FUNCTION
  auto &operator()(int entity, int node, int component, Mesh_Topology t) const {
    if (t != Vertex) {
      Kokkos::printf("%d is not a support topology\n", t);
      assert(false);
    }
    return vtxField(entity, node, component);
  }
};
 
template <typename ExecutionSpace,
          template <typename...>
          typename Controller = MeshField::KokkosController,
          typename DataType, size_t order, size_t dim, size_t numComp>
auto CreateLagrangeField(const MeshInfo &meshInfo) {
  static_assert((std::is_same_v<Real4, DataType> == true ||
                 std::is_same_v<Real8, DataType> == true),
                "CreateLagrangeField only supports single and double precision "
                "floating point fields\n");
  static_assert(
      (order == 1 || order == 2),
      "CreateLagrangeField only supports linear and quadratic fields\n");
  static_assert((dim == 1 || dim == 2 || dim == 3),
                "CreateLagrangeField only supports 1d, 2d, and 3d meshes\n");
  using MemorySpace = typename ExecutionSpace::memory_space;
  if constexpr (order == 1 && (dim == 1 || dim == 2 || dim == 3)) {
    if (meshInfo.numVtx <= 0) {
      fail("mesh has no vertices\n");
    }
#ifdef MESHFIELDS_ENABLE_CABANA
    using Ctrlr = std::conditional_t<
        std::is_same_v<
            Controller<ExecutionSpace, MemorySpace, DataType>,
            MeshField::CabanaController<ExecutionSpace, MemorySpace, DataType>>,
        Controller<ExecutionSpace, MemorySpace, DataType[1][numComp]>,
        Controller<MemorySpace, ExecutionSpace, DataType ***>>;
    // 1 dof with 1 component per vtx
    auto createController = [](auto numVtx) {
      if constexpr (std::is_same_v<
                        Controller<ExecutionSpace, MemorySpace, DataType>,
                        MeshField::CabanaController<ExecutionSpace, MemorySpace,
                                                    DataType>>) {
        return Ctrlr({numVtx});
      } else {
        return Ctrlr({/*field 0*/ numVtx, 1, numComp});
      }
    };
    Ctrlr kk_ctrl = createController(meshInfo.numVtx);
#else
    using Ctrlr = Controller<MemorySpace, ExecutionSpace, DataType ***>;
    Ctrlr kk_ctrl({/*field 0*/ meshInfo.numVtx, 1, numComp});
#endif
    auto vtxField = MeshField::makeField<Ctrlr, 0>(kk_ctrl);
    using LA = LinearAccessor<decltype(vtxField)>;
    // clang-format off
    using LinearLagrangeShapeField = std::conditional_t<
        dim == 3,
        ShapeField<numComp, LinearTetrahedronShape, LA>,
        ShapeField<numComp, LinearTriangleShape, LA>>;
    // clang-format on
    LinearLagrangeShapeField llsf(meshInfo, {vtxField});
    return FieldWithController<Ctrlr, LinearLagrangeShapeField>{kk_ctrl, llsf};
  } else if constexpr (order == 2 && (dim == 2 || dim == 3)) {
    if (meshInfo.numVtx <= 0) {
      fail("mesh has no vertices\n");
    }
    if (meshInfo.numEdge <= 0) {
      fail("mesh has no edges\n");
    }
#ifdef MESHFIELDS_ENABLE_CABANA
    using Ctrlr = std::conditional_t<
        std::is_same_v<
            Controller<ExecutionSpace, MemorySpace, DataType>,
            MeshField::CabanaController<ExecutionSpace, MemorySpace, DataType>>,
        Controller<ExecutionSpace, MemorySpace, DataType[1][numComp],
                   DataType[1][numComp]>,
        Controller<MemorySpace, ExecutionSpace, DataType ***, DataType ***>>;
    // 1 dof with 1 comp per vtx/edge
    auto createController = [](auto numVtx, auto numEdge) {
      if constexpr (std::is_same_v<
                        Controller<ExecutionSpace, MemorySpace, DataType>,
                        MeshField::CabanaController<ExecutionSpace, MemorySpace,
                                                    DataType>>) {
        return Ctrlr({numVtx, numEdge});
      } else {
        return Ctrlr({/*field 0*/ numVtx, 1, numComp,
                      /*field 1*/ numEdge, 1, numComp});
      }
    };
    Ctrlr kk_ctrl = createController(meshInfo.numVtx, meshInfo.numEdge);
#else
    using Ctrlr =
        Controller<MemorySpace, ExecutionSpace, DataType ***, DataType ***>;
    Ctrlr kk_ctrl({/*field 0*/ meshInfo.numVtx, 1, numComp,
                   /*field 1*/ meshInfo.numEdge, 1, numComp});
#endif
    auto vtxField = MeshField::makeField<Ctrlr, 0>(kk_ctrl);
    auto edgeField = MeshField::makeField<Ctrlr, 1>(kk_ctrl);
    using QA = QuadraticAccessor<decltype(vtxField), decltype(edgeField)>;
    // clang-format off
    using QuadraticLagrangeShapeField = std::conditional_t<
        dim == 3,
        ShapeField<numComp, QuadraticTetrahedronShape, QA>,
        ShapeField<numComp, QuadraticTriangleShape, QA>>;
    // clang-format on
    QuadraticLagrangeShapeField qlsf(meshInfo, {vtxField, edgeField});
    return FieldWithController<Ctrlr, QuadraticLagrangeShapeField>{kk_ctrl,
                                                                    qlsf};
  } else {
    fail("CreateLagrangeField does not support the specified "
         "combination of order %d and dimension %d.\n",
         order, dim);
    return FieldWithController<int, std::nullptr_t>{}; // silence compiler warning
  }
};
 
template <typename ExecutionSpace,
          template <typename...>
          typename Controller = MeshField::KokkosController,
          size_t dim>
auto CreateCoordinateField(const MeshInfo &meshInfo) {
  // current coordinate fields are assumed to be linear over the elements
  return CreateLagrangeField<ExecutionSpace, Controller, Real, 1, dim, dim>(meshInfo);
}
 
} // namespace MeshField
 
#endif