KokkosFESystem.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 "KokkosSystem.h"
#include "KokkosAssembly.h"
 
class MooseMesh;
class SystemBase;
 
namespace Moose::Kokkos
{
 
class NodalBCBase;
 
class FESystem : public System, public AssemblyHolder
{
public:
  FESystem(SystemBase & system);
 
  FESystem(System & base, SystemBase & system);
 
#ifdef MOOSE_KOKKOS_SCOPE
 
  void reinit();
 
  KOKKOS_FUNCTION auto getLocalConstraints(dof_id_type dof) const
  {
    return _local_constraints.size() ? _local_constraints[dof]
                                     : decltype(_local_constraints[dof])(nullptr, {});
  }
 
  KOKKOS_FUNCTION const Array<unsigned int> & getFieldCoupling(unsigned int var) const
  {
    return _field_coupling[var];
  }
 
  KOKKOS_FUNCTION const Array<unsigned int> & getScalarCoupling(unsigned int var) const
  {
    return _scalar_coupling[var];
  }
 
  KOKKOS_FUNCTION bool hasNodalBCMatrixTag(dof_id_type dof, TagID tag) const
  {
    return _nbc_matrix_tag_dof[tag].isAlloc() && _nbc_matrix_tag_dof[tag][dof];
  }
 
  KOKKOS_FUNCTION unsigned int getFETypeID(unsigned int var) const { return _var_fe_types[var]; }
 
  KOKKOS_FUNCTION dof_id_type getNodeLocalDofIndex(ContiguousNodeID node,
                                                   unsigned int i,
                                                   unsigned int var) const
  {
    return _local_node_dof_index[var][node] + i;
  }
 
  KOKKOS_FUNCTION dof_id_type getNodeGlobalDofIndex(ContiguousNodeID node,
                                                    unsigned int i,
                                                    unsigned int var) const
  {
    return _local_to_global_dof_index[getNodeLocalDofIndex(node, i, var)];
  }
 
  KOKKOS_FUNCTION bool isNodalDefined(ContiguousNodeID node, unsigned int var) const
  {
    return _local_node_dof_index[var][node] != libMesh::DofObject::invalid_id;
  }
 
  KOKKOS_FUNCTION ADReal getVectorDofADValue(const dof_id_type dof,
                                             const TagID tag,
                                             const Real seed) const;
  KOKKOS_FUNCTION Real & getVectorQpValue(const ElementInfo info,
                                          const dof_id_type qp,
                                          const unsigned int var,
                                          const TagID tag) const
  {
    return _qp_solutions[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION ADReal getVectorQpADValue(const ElementInfo info,
                                            const dof_id_type offset,
                                            const dof_id_type qp,
                                            const unsigned int var,
                                            const TagID tag,
                                            const Real seed) const;
  KOKKOS_FUNCTION Real3 & getVectorQpGrad(const ElementInfo info,
                                          const dof_id_type qp,
                                          const unsigned int var,
                                          const TagID tag) const
  {
    return _qp_solutions_grad[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION Real3 & getVectorQpVectorValue(const ElementInfo info,
                                                 const dof_id_type qp,
                                                 const unsigned int var,
                                                 const TagID tag) const
  {
    return _qp_vector_solutions[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION Real33 & getVectorQpVectorGrad(const ElementInfo info,
                                                 const dof_id_type qp,
                                                 const unsigned int var,
                                                 const TagID tag) const
  {
    return _qp_vector_solutions_grad[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION Real3 & getVectorQpVectorCurl(const ElementInfo info,
                                                const dof_id_type qp,
                                                const unsigned int var,
                                                const TagID tag) const
  {
    return _qp_vector_solutions_curl[tag](info.subdomain, var)[qp];
  }
  KOKKOS_FUNCTION ADReal3 getVectorQpADGrad(const ElementInfo info,
                                            const Real33 jacobian,
                                            const dof_id_type offset,
                                            const dof_id_type qp,
                                            const unsigned int var,
                                            const TagID tag,
                                            const Real seed) const;
  KOKKOS_FUNCTION Real getVectorQpValueFace(const ElementInfo info,
                                            const unsigned int side,
                                            const unsigned int qp,
                                            const unsigned int var,
                                            const TagID tag) const;
  KOKKOS_FUNCTION Real3 getVectorQpVectorValueFace(const ElementInfo info,
                                                   const unsigned int side,
                                                   const unsigned int qp,
                                                   const unsigned int var,
                                                   const TagID tag) const;
  KOKKOS_FUNCTION ADReal getVectorQpADValueFace(const ElementInfo info,
                                                const unsigned int side,
                                                const unsigned int qp,
                                                const unsigned int var,
                                                const TagID tag,
                                                const Real seed) const;
  KOKKOS_FUNCTION Real3 getVectorQpGradFace(const ElementInfo info,
                                            const unsigned int side,
                                            const Real33 jacobian,
                                            const unsigned int qp,
                                            const unsigned int var,
                                            const TagID tag) const;
  KOKKOS_FUNCTION Real33 getVectorQpVectorGradFace(const ElementInfo info,
                                                   const unsigned int side,
                                                   const Real33 jacobian,
                                                   const unsigned int qp,
                                                   const unsigned int var,
                                                   const TagID tag) const;
 
  KOKKOS_FUNCTION ADReal3 getVectorQpADGradFace(const ElementInfo info,
                                                const unsigned int side,
                                                const Real33 jacobian,
                                                const unsigned int qp,
                                                const unsigned int var,
                                                const TagID tag,
                                                const Real seed) const;
 
  KOKKOS_FUNCTION void operator()(const ThreadID tid) const;
#endif
 
private:
  void setupVariables();
 
  void setupDofs();
 
  void setupCoupling();
 
  void setupNodalBCDofs();
 
  void getNodalBCDofs(const NodalBCBase * nbc, Array<bool> & dofs);
 
  Thread<> _thread;
 
  Array<Array2D<Array<Real>>> _qp_solutions;
  Array<Array2D<Array<Real3>>> _qp_solutions_grad;
  Array<Array2D<Array<Real3>>> _qp_vector_solutions;
  Array<Array2D<Array<Real33>>> _qp_vector_solutions_grad;
  Array<Array2D<Array<Real3>>> _qp_vector_solutions_curl;
 
  Array<Array<dof_id_type>> _local_node_dof_index;
 
  Array<unsigned int> _var_fe_types;
 
  Array<bool> _var_is_vector;
 
  Array<Array<unsigned int>> _field_coupling;
  Array<Array<unsigned int>> _scalar_coupling;
 
  Array<Array<bool>> _nbc_matrix_tag_dof;
 
  JaggedArray<Pair<dof_id_type, Real>, 1, 1> _local_constraints;
};
 
#ifdef MOOSE_KOKKOS_SCOPE
 
KOKKOS_FUNCTION inline ADReal
FESystem::getVectorDofADValue(const dof_id_type dof, TagID tag, const Real seed) const
{
  ADReal value = _vectors[tag][dof];
 
  if (seed != 0)
    value.derivatives().insert(dof) = seed;
 
  return value;
}
 
KOKKOS_FUNCTION inline ADReal
FESystem::getVectorQpADValue(const ElementInfo info,
                             const dof_id_type offset,
                             const dof_id_type qp,
                             const unsigned int var,
                             const TagID tag,
                             const Real seed) const
{
  ADReal value = 0;
 
  if (seed == 0)
    value = getVectorQpValue(info, offset + qp, var, tag);
  else
  {
    auto fe = _var_fe_types[var];
    auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
    auto & phi = kokkosAssembly().getPhi(info.subdomain, info.type, fe);
 
    for (unsigned int i = 0; i < n_dofs; ++i)
      value += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * phi(i, qp);
  }
 
  return value;
}
 
KOKKOS_FUNCTION inline ADReal3
FESystem::getVectorQpADGrad(const ElementInfo info,
                            const Real33 jacobian,
                            const dof_id_type offset,
                            const dof_id_type qp,
                            const unsigned int var,
                            const TagID tag,
                            const Real seed) const
{
  ADReal3 grad;
 
  if (seed == 0)
    grad = getVectorQpGrad(info, offset + qp, var, tag);
  else
  {
    auto fe = _var_fe_types[var];
    auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
    auto & grad_phi = kokkosAssembly().getGradPhi(info.subdomain, info.type, fe);
 
    for (unsigned int i = 0; i < n_dofs; ++i)
      grad +=
          getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * grad_phi(i, qp);
 
    grad = jacobian * grad;
  }
 
  return grad;
}
 
KOKKOS_FUNCTION inline Real
FESystem::getVectorQpValueFace(const ElementInfo info,
                               const unsigned int side,
                               const unsigned int qp,
                               const unsigned int var,
                               const TagID tag) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);
 
  Real value = 0;
 
  for (unsigned int i = 0; i < n_dofs; ++i)
    value += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * phi(i, qp);
 
  return value;
}
 
KOKKOS_FUNCTION inline Real3
FESystem::getVectorQpVectorValueFace(const ElementInfo info,
                                     const unsigned int side,
                                     const unsigned int qp,
                                     const unsigned int var,
                                     const TagID tag) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & phi = kokkosAssembly().getVectorPhiFace(info.subdomain, info.type, fe)(side);
 
  Real3 value = 0;
 
  for (unsigned int i = 0; i < n_dofs; ++i)
    value += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * phi(i, qp);
 
  return value;
}
 
KOKKOS_FUNCTION inline ADReal
FESystem::getVectorQpADValueFace(const ElementInfo info,
                                 const unsigned int side,
                                 const unsigned int qp,
                                 const unsigned int var,
                                 const TagID tag,
                                 const Real seed) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & phi = kokkosAssembly().getPhiFace(info.subdomain, info.type, fe)(side);
 
  ADReal value = 0;
 
  for (unsigned int i = 0; i < n_dofs; ++i)
    value += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * phi(i, qp);
 
  return value;
}
 
KOKKOS_FUNCTION inline Real3
FESystem::getVectorQpGradFace(const ElementInfo info,
                              const unsigned int side,
                              const Real33 jacobian,
                              const unsigned int qp,
                              const unsigned int var,
                              const TagID tag) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);
 
  Real3 grad = 0;
 
  for (unsigned int i = 0; i < n_dofs; ++i)
    grad += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * grad_phi(i, qp);
 
  grad = jacobian * grad;
 
  return grad;
}
 
KOKKOS_FUNCTION inline Real33
FESystem::getVectorQpVectorGradFace(const ElementInfo info,
                                    const unsigned int side,
                                    const Real33 jacobian,
                                    const unsigned int qp,
                                    const unsigned int var,
                                    const TagID tag) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & grad_phi = kokkosAssembly().getVectorGradPhiFace(info.subdomain, info.type, fe)(side);
 
  Real33 grad = 0;
 
  for (unsigned int i = 0; i < n_dofs; ++i)
    grad += getVectorDofValue(getElemLocalDofIndex(info.id, i, var), tag) * grad_phi(i, qp);
 
  grad = grad * jacobian.transpose();
 
  return grad;
}
 
KOKKOS_FUNCTION inline ADReal3
FESystem::getVectorQpADGradFace(const ElementInfo info,
                                const unsigned int side,
                                const Real33 jacobian,
                                const unsigned int qp,
                                const unsigned int var,
                                const TagID tag,
                                const Real seed) const
{
  auto fe = _var_fe_types[var];
  auto n_dofs = kokkosAssembly().getNumDofs(info.type, fe);
  auto & grad_phi = kokkosAssembly().getGradPhiFace(info.subdomain, info.type, fe)(side);
 
  ADReal3 grad = ADReal(0);
 
  for (unsigned int i = 0; i < n_dofs; ++i)
    grad += getVectorDofADValue(getElemLocalDofIndex(info.id, i, var), tag, seed) * grad_phi(i, qp);
 
  grad = jacobian * grad;
 
  return grad;
}
#endif
 
MakeSystemHolder(FESystem);
} // namespace Moose::Kokkos