text: // Copyright (c) 2010-2021, Lawrence Livermore National Security, LLC. Produced
  // at the Lawrence Livermore National Laboratory. All Rights reserved. See files
  // LICENSE and NOTICE for details. LLNL-CODE-806117.
  //
  // This file is part of the MFEM library. For more information and source code
  // availability visit https://mfem.org.
  //
  // MFEM is free software; you can redistribute it and/or modify it under the
  // terms of the BSD-3 license. We welcome feedback and contributions, see file
  // CONTRIBUTING.md for details.
  
  #ifndef MFEM_TEMPLATE_FESPACE
  #define MFEM_TEMPLATE_FESPACE
  
  #include "../config/tconfig.hpp"
  #include "../general/tassign.hpp"
  #include "../linalg/ttensor.hpp"
  #include "tfe.hpp" // for TFiniteElementSpace_simple
  #include "fespace.hpp"
  
  namespace mfem
  {
  
  // Templated finite element space classes, cf. fespace.?pp and fe_coll.?pp
  
  // Index types
  
  // IndexType must define:
  // - constructor IndexType(const FE &fe, const FiniteElementSpace &fes)
  // - copy constructor
  // - void SetElement(int elem_idx)
  // - int map(int loc_dof_idx, int elem_offset) const --> glob_dof_idx, for
  //   single component; elem_offset is relative to the currently set element.
  
  // Index type based on an array listing the dofs for each element where all
  // elements are assumed to have the same number of dofs. Such an array is
  // constructed from the J array of an element-to-dof Table with optional local
  // renumbering to ensure tensor-product local dof ordering when needed.
  template <typename FE>
  class ElementDofIndexer
  {
  protected:
     const int *el_dof_list, *loc_dof_list;
     bool own_list;
  
  public:
     typedef FE FE_type;
  
     ElementDofIndexer(const FE &fe, const FiniteElementSpace &fes)
     {
        const Array<int> *loc_dof_map = fe.GetDofMap();
        // fes.BuildElementToDofTable();
        const Table &el_dof = fes.GetElementToDofTable();
        MFEM_ASSERT(el_dof.Size_of_connections() == el_dof.Size() * FE::dofs,
                    "the element-to-dof Table is not compatible with this FE!");
        int num_dofs = el_dof.Size() * FE::dofs;
        if (!loc_dof_map)
        {
           // no local dof reordering
           el_dof_list = el_dof.GetJ();
           own_list = false;
        }
        else
        {
           // reorder the local dofs according to loc_dof_map
           int *el_dof_list_ = new int[num_dofs];
           const int *loc_dof_map_ = loc_dof_map->GetData();
           for (int i = 0; i < el_dof.Size(); i++)
           {
              MFEM_ASSERT(el_dof.RowSize(i) == FE::dofs,
                          "incompatible element-to-dof Table!");
              for (int j = 0; j < FE::dofs; j++)
              {
                 el_dof_list_[j+FE::dofs*i] =
                    el_dof.GetJ()[loc_dof_map_[j]+FE::dofs*i];
              }
           }
           el_dof_list = el_dof_list_;
           own_list = true;
        }
        loc_dof_list = el_dof_list; // point to element 0
     }
  
     // Shallow copy constructor
     inline MFEM_ALWAYS_INLINE
     ElementDofIndexer(const ElementDofIndexer &orig)
        : el_dof_list(orig.el_dof_list),
          loc_dof_list(orig.loc_dof_list),
          own_list(false)
     { }
  
     inline MFEM_ALWAYS_INLINE
     ~ElementDofIndexer() { if (own_list) { delete [] el_dof_list; } }
  
     inline MFEM_ALWAYS_INLINE
     void SetElement(int elem_idx)
     {
        loc_dof_list = el_dof_list + elem_idx * FE::dofs;
     }
  
     inline MFEM_ALWAYS_INLINE
     int map(int loc_dof_idx, int elem_offset) const
     {
        return loc_dof_list[loc_dof_idx + elem_offset * FE::dofs];
     }
  };
  
  
  // Simple template Finite Element Space, built using an IndexType. For a
  // description of the requirements on IndexType, see above.
  template <typename FE, typename IndexType>
  class TFiniteElementSpace_simple
  {
  public:
     typedef FE        FE_type;
     typedef IndexType index_type;
     static const int dofs = FE::dofs;
  
  protected:
     index_type ind;
     int num_elems, remain_elems;
  
  public:
     TFiniteElementSpace_simple(const FE &fe, const FiniteElementSpace &fes)
        : ind(fe, fes), num_elems(fes.GetNE()), remain_elems(num_elems) { }
  
     // default copy constructor
  
     int GetNE() const { return num_elems; }
  
     void SetElement(int el) { ind.SetElement(el); remain_elems = num_elems-el; }
  
  #if 0
     // Multi-element Extract:
     // Extract dofs for multiple elements starting with the current element.
     // The number of elements to extract is given by the second dimension of
     // dof_layout_t: dof_layout is (DOFS x NumElems).
     template <AssignOp::Type Op, typename glob_dof_data_t,
               typename dof_layout_t, typename dof_data_t>
     inline MFEM_ALWAYS_INLINE
     void Extract(const glob_dof_data_t &glob_dof_data,
                  const dof_layout_t    &dof_layout,
                  dof_data_t            &dof_data) const
     {
        const int SS = sizeof(dof_data[0])/sizeof(dof_data[0][0]);
        const int NE = dof_layout_t::dim_2;
        MFEM_STATIC_ASSERT(FE::dofs == dof_layout_t::dim_1,
                           "invalid number of dofs");
        for (int j = 0; j < NE; j++)
        {
           for (int i = 0; i < FE::dofs; i++)
           {
              for (int s = 0; s < SS; s++)
              {
                 Assign<Op>(dof_data[dof_layout.ind(i,j)][s],
                            glob_dof_data[ind.map(i,s+SS*j)]);
              }
           }
        }
     }
  
     template <typename glob_dof_data_t,
               typename dof_layout_t, typename dof_data_t>
     inline MFEM_ALWAYS_INLINE
     void Extract(const glob_dof_data_t &glob_dof_data,
                  const dof_layout_t    &dof_layout,
                  dof_data_t            &dof_data) const
     {
        Extract<AssignOp::Set>(glob_dof_data, dof_layout, dof_data);
     }
  
     // Multi-element assemble.
     template <AssignOp::Type Op,
               typename dof_layout_t, typename dof_data_t,
               typename glob_dof_data_t>
     inline MFEM_ALWAYS_INLINE
     void Assemble(const dof_layout_t &dof_layout,
                   const dof_data_t   &dof_data,
                   glob_dof_data_t    &glob_dof_data) const
     {
        const int SS = sizeof(dof_data[0])/sizeof(dof_data[0][0]);
        const int NE = dof_layout_t::dim_2;
        MFEM_STATIC_ASSERT(FE::dofs == dof_layout_t::dim_1,
                           "invalid number of dofs");
        for (int j = 0; j < NE; j++)
        {
           for (int i = 0; i < FE::dofs; i++)
           {
              for (int s = 0; s < SS; s++)
              {
                 Assign<Op>(glob_dof_data[ind.map(i,s+SS*j)],
                            dof_data[dof_layout.ind(i,j)][s]);
              }
           }
        }
     }
  
     template <typename dof_layout_t, typename dof_data_t,
               typename glob_dof_data_t>
     inline MFEM_ALWAYS_INLINE
     void Assemble(const dof_layout_t &dof_layout,
                   const dof_data_t   &dof_data,
                   glob_dof_data_t    &glob_dof_data) const
     {
        Assemble<AssignOp::Add>(dof_layout, dof_data, glob_dof_data);
     }
  #endif
  
     // Multi-element VectorExtract: vdof_layout is (DOFS x NumComp x NumElems).
     template <AssignOp::Type Op,
               typename vec_layout_t, typename glob_vdof_data_t,
               typename vdof_layout_t, typename vdof_data_t>
     inline MFEM_ALWAYS_INLINE
     void VectorExtract(const vec_layout_t     &vl,
                        const glob_vdof_data_t &glob_vdof_data,
                        const vdof_layout_t    &vdof_layout,
                        vdof_data_t            &vdof_data) const
     {
        const int SS = sizeof(vdof_data[0])/sizeof(vdof_data[0][0]);
        const int NC = vdof_layout_t::dim_2;
        const int NE = vdof_layout_t::dim_3;
        MFEM_STATIC_ASSERT(FE::dofs == vdof_layout_t::dim_1,
                           "invalid number of dofs");
        MFEM_ASSERT(NC == vl.NumComponents(), "invalid number of components");
        const int TE = std::min(SS*NE, remain_elems);
        // const int TE = SS*NE;
        for (int k = 0; k < NC; k++)
        {
  #if 0
           for (int j = 0; j < NE; j++)
           {
              for (int i = 0; i < FE::dofs; i++)
              {
                 for (int s = 0; s < SS; s++)
                 {
                    Assign<Op>(vdof_data[vdof_layout.ind(i,k,j)][s],
                               glob_vdof_data[vl.ind(ind.map(i,s+SS*j), k)]);
                 }
              }
           }
  #else
           for (int js = 0; js < TE; js++)
           {
              for (int i = 0; i < FE::dofs; i++)
              {
                 const int s = js % SS, j = js / SS;
                 Assign<Op>(vdof_data[vdof_layout.ind(i,k,j)][s],
                            glob_vdof_data[vl.ind(ind.map(i,js), k)]);
              }
           }
  #endif
        }
     }
  
     template <typename vec_layout_t, typename glob_vdof_data_t,
               typename vdof_layout_t, typename vdof_data_t>
     inline MFEM_ALWAYS_INLINE
     void VectorExtract(const vec_layout_t     &vl,
                        const glob_vdof_data_t &glob_vdof_data,
                        const vdof_layout_t    &vdof_layout,
                        vdof_data_t            &vdof_data) const
     {
        VectorExtract<AssignOp::Set>(vl, glob_vdof_data, vdof_layout, vdof_data);
     }
  
     // Multi-element VectorAssemble: vdof_layout is (DOFS x NumComp x NumElems).
     template <AssignOp::Type Op,
               typename vdof_layout_t, typename vdof_data_t,
               typename vec_layout_t, typename glob_vdof_data_t>
     inline MFEM_ALWAYS_INLINE
     void VectorAssemble(const vdof_layout_t &vdof_layout,
                         const vdof_data_t   &vdof_data,
                         const vec_layout_t  &vl,
                         glob_vdof_data_t    &glob_vdof_data) const
     {
        const int SS = sizeof(vdof_data[0])/sizeof(vdof_data[0][0]);
        const int NC = vdof_layout_t::dim_2;
        const int NE = vdof_layout_t::dim_3;
        MFEM_STATIC_ASSERT(FE::dofs == vdof_layout_t::dim_1,
                           "invalid number of dofs");
        MFEM_ASSERT(NC == vl.NumComponents(), "invalid number of components");
        const int TE = std::min(SS*NE, remain_elems);
        // const int TE = SS*NE;
        for (int k = 0; k < NC; k++)
        {
  #if 0
           for (int j = 0; j < NE; j++)
           {
              for (int i = 0; i < FE::dofs; i++)
              {
                 for (int s = 0; s < SS; s++)
                 {
                    Assign<Op>(glob_vdof_data[vl.ind(ind.map(i,s+SS*j), k)],
                               vdof_data[vdof_layout.ind(i,k,j)][s]);
                 }
              }
           }
  #else
           for (int js = 0; js < TE; js++)
           {
              for (int i = 0; i < FE::dofs; i++)
              {
                 const int s = js % SS, j = js / SS;
                 Assign<Op>(glob_vdof_data[vl.ind(ind.map(i,js), k)],
                            vdof_data[vdof_layout.ind(i,k,j)][s]);
              }
           }
  #endif
        }
     }
  
     template <typename vdof_layout_t, typename vdof_data_t,
               typename vec_layout_t, typename glob_vdof_data_t>
     inline MFEM_ALWAYS_INLINE
     void VectorAssemble(const vdof_layout_t &vdof_layout,
                         const vdof_data_t   &vdof_data,
                         const vec_layout_t  &vl,
                         glob_vdof_data_t    &glob_vdof_data) const
     {
        VectorAssemble<AssignOp::Add>(vdof_layout, vdof_data, vl, glob_vdof_data);
     }
  
     // Extract a static number of consecutive components; vdof_layout is
     // (dofs x NC x NE), where NC is the number of components to extract. It is
     // assumed that: first_comp + NC <= vl.NumComponents().
     template <typename vdof_layout_t, typename vdof_data_t,
               typename vec_layout_t, typename glob_vdof_data_t>
     inline MFEM_ALWAYS_INLINE
     void ExtractComponents(int                     first_comp,
                            const vec_layout_t     &vl,
                            const glob_vdof_data_t &glob_vdof_data,
                            const vdof_layout_t    &vdof_layout,
                            vdof_data_t            &vdof_data) const
     {
        const int SS = sizeof(vdof_data[0])/sizeof(vdof_data[0][0]);
        const int NC = vdof_layout_t::dim_2;
        const int NE = vdof_layout_t::dim_3;
        const int TE = std::min(SS*NE, remain_elems);
        MFEM_STATIC_ASSERT(FE::dofs == vdof_layout_t::dim_1,
                           "invalid number of dofs");
        MFEM_ASSERT(first_comp + NC <= vl.NumComponents(),
                    "invalid number of components");
        for (int k = 0; k < NC; k++)
        {
           for (int js = 0; js < TE; js++)
           {
              for (int i = 0; i < FE::dofs; i++)
              {
                 const int s = js % SS, j = js / SS;
                 Assign<AssignOp::Set>(
                    vdof_data[vdof_layout.ind(i,k,j)][s],
                    glob_vdof_data[vl.ind(ind.map(i,js), first_comp+k)]);
              }
           }
        }
     }
  
     // Assemble a static number of consecutive components; vdof_layout is
     // (dofs x NC x NE), where NC is the number of components to add. It is
     // assumed that: first_comp + NC <= vl.NumComponents().
     template <typename vdof_layout_t, typename vdof_data_t,
               typename vec_layout_t, typename glob_vdof_data_t>
     inline MFEM_ALWAYS_INLINE
     void AssembleComponents(int                  first_comp,
                             const vdof_layout_t &vdof_layout,
                             const vdof_data_t   &vdof_data,
                             const vec_layout_t  &vl,
                             glob_vdof_data_t    &glob_vdof_data) const
     {
        const int SS = sizeof(vdof_data[0])/sizeof(vdof_data[0][0]);
        const int NC = vdof_layout_t::dim_2;
        const int NE = vdof_layout_t::dim_3;
        const int TE = std::min(SS*NE, remain_elems);
        MFEM_STATIC_ASSERT(FE::dofs == vdof_layout_t::dim_1,
                           "invalid number of dofs");
        MFEM_ASSERT(first_comp + NC <= vl.NumComponents(),
                    "invalid number of components");
        for (int k = 0; k < NC; k++)
        {
           for (int js = 0; js < TE; js++)
           {
              for (int i = 0; i < FE::dofs; i++)
              {
                 const int s = js % SS, j = js / SS;
                 Assign<AssignOp::Add>(
                    glob_vdof_data[vl.ind(ind.map(i,js), first_comp+k)],
                    vdof_data[vdof_layout.ind(i,k,j)][s]);
              }
           }
        }
     }
  
     template <typename vcomplex_t>
     void Assemble(const TMatrix<FE::dofs,FE::dofs,vcomplex_t> &m,
                   SparseMatrix &M) const
     {
        const int SS = sizeof(m[0])/sizeof(m[0][0]);
        const int TE = std::min(SS, remain_elems);
        MFEM_FLOPS_ADD(FE::dofs*FE::dofs);
        for (int s = 0; s < TE; s++)
        {
           for (int i = 0; i < FE::dofs; i++)
           {
              M.SetColPtr(ind.map(i,s));
              for (int j = 0; j < FE::dofs; j++)
              {
                 M._Add_(ind.map(j,s), m(i,j)[s]);
              }
              M.ClearColPtr();
           }
        }
     }
  
     template <typename vec_layout_t, typename vcomplex_t>
     void AssembleBlock(int block_i, int block_j, const vec_layout_t &vl,
                        const TMatrix<FE::dofs,FE::dofs,vcomplex_t> &m,
                        SparseMatrix &M) const
     {
        const int SS = sizeof(m[0])/sizeof(m[0][0]);
        const int TE = std::min(SS, remain_elems);
        MFEM_FLOPS_ADD(FE::dofs*FE::dofs);
        for (int s = 0; s < TE; s++)
        {
           for (int i = 0; i < FE::dofs; i++)
           {
              M.SetColPtr(vl.ind(ind.map(i,s), block_i));
              for (int j = 0; j < FE::dofs; j++)
              {
                 M._Add_(vl.ind(ind.map(j,s), block_j), m(i,j)[s]);
              }
              M.ClearColPtr();
           }
        }
     }
  };
  
  // H1 Finite Element Space
  
  template <typename FE>
  class H1_FiniteElementSpace
     : public TFiniteElementSpace_simple<FE,ElementDofIndexer<FE> >
  {
  public:
     typedef FE FE_type;
     typedef TFiniteElementSpace_simple<FE,ElementDofIndexer<FE> > base_class;
  
     H1_FiniteElementSpace(const FE &fe, const FiniteElementSpace &fes)
        : base_class(fe, fes)
     { }
  
     // default copy constructor
  
     static bool Matches(const FiniteElementSpace &fes)
     {
        const FiniteElementCollection *fec = fes.FEColl();
        const H1_FECollection *h1_fec =
           dynamic_cast<const H1_FECollection *>(fec);
        if (!h1_fec) { return false; }
        const FiniteElement *fe = h1_fec->FiniteElementForGeometry(FE_type::geom);
        if (fe->GetOrder() != FE_type::degree) { return false; }
        return true;
     }
  
     template <typename vec_layout_t>
     static bool VectorMatches(const FiniteElementSpace &fes)
     {
        return Matches(fes) && vec_layout_t::Matches(fes);
     }
  };
  
  
  // Simple index type for DG spaces, where the map method is given by:
  // glob_dof_idx = loc_dof_idx + elem_idx * num_dofs.
  template <typename FE>
  class DGIndexer
  {
  protected:
     int offset;
  
  public:
     typedef FE FE_type;
  
     DGIndexer(const FE &fe, const FiniteElementSpace &fes)
     {
        MFEM_ASSERT(fes.GetNDofs() == fes.GetNE() * FE::dofs,
                    "the FE space is not compatible with this FE!");
        offset = 0;
     }
  
     // default copy constructor
  
     inline MFEM_ALWAYS_INLINE
     void SetElement(int elem_idx)
     {
        offset = FE::dofs * elem_idx;
     }
  
     inline MFEM_ALWAYS_INLINE
     int map(int loc_dof_idx, int elem_offset) const
     {
        return offset + loc_dof_idx + elem_offset * FE::dofs;
     }
  };
  
  
  // L2 Finite Element Space
  
  template <typename FE>
  class L2_FiniteElementSpace
     : public TFiniteElementSpace_simple<FE,DGIndexer<FE> >
  {
  public:
     typedef FE FE_type;
     typedef TFiniteElementSpace_simple<FE,DGIndexer<FE> > base_class;
  
     L2_FiniteElementSpace(const FE &fe, const FiniteElementSpace &fes)
        : base_class(fe, fes)
     { }
  
     // default copy constructor
  
     static bool Matches(const FiniteElementSpace &fes)
     {
        const FiniteElementCollection *fec = fes.FEColl();
        const L2_FECollection *l2_fec =
           dynamic_cast<const L2_FECollection *>(fec);
        if (!l2_fec) { return false; }
        const FiniteElement *fe = l2_fec->FiniteElementForGeometry(FE_type::geom);
        if (fe->GetOrder() != FE_type::degree) { return false; }
        return true;
     }
  
     template <typename vec_layout_t>
     static bool VectorMatches(const FiniteElementSpace &fes)
     {
        return Matches(fes) && vec_layout_t::Matches(fes);
     }
  };
  
  } // namespace mfem
  
  #endif // MFEM_TEMPLATE_FESPACE
  