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;     // entDim = 0
  int numEdge;    // entDim = 1
  int numTri;     // entDim = 2
  int numQuad;    // entDim = 2
  int numTet;     // entDim = 3
  int numHex;     // entDim = 3
  int numPrism;   // entDim = 3
  int numPyramid; // entDim = 3
  int dim;
};
 
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) {
  if (meshInfo.numVtx <= 0) {
    fail("mesh has no vertices\n");
  }
  using DataType = Real;
  using MemorySpace = typename ExecutionSpace::memory_space;
  const int numComp = meshInfo.dim;
#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][dim]>,
      Controller<MemorySpace, ExecutionSpace, DataType ***>>;
  auto createController = [](const int numComp, 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(numComp, 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)>;
  using LinearLagrangeShapeField =
      ShapeField<dim, LinearTriangleShape, LA>;
  LinearLagrangeShapeField llsf(meshInfo, {vtxField});
  return FieldWithController<Ctrlr, LinearLagrangeShapeField>{kk_ctrl, llsf};
};
 
} // namespace MeshField
 
#endif