KokkosVariableValue.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* 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 "KokkosDatum.h"

#include "MooseVariableBase.h"

namespace Moose
{
namespace Kokkos
{

class VariablePhiValue
{
public:
  KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.jfe();

    return side == libMesh::invalid_uint
               ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
               : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
  }
};

class VariablePhiGradient
{
public:
  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.jfe();

    return datum.J(qp) *
           (side == libMesh::invalid_uint
                ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
                : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
  }
};

class VariableTestValue
{
public:
  KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.ife();

    return side == libMesh::invalid_uint
               ? datum.assembly().getPhi(elem.subdomain, elem.type, fe)(i, qp)
               : datum.assembly().getPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);
  }
};

class VariableTestGradient
{
public:
  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = datum.ife();

    return datum.J(qp) *
           (side == libMesh::invalid_uint
                ? datum.assembly().getGradPhi(elem.subdomain, elem.type, fe)(i, qp)
                : datum.assembly().getGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp));
  }
};

class VariableValue
{
public:
  VariableValue() = default;
  VariableValue(Variable var, bool dof = false) : _var(var), _dof(dof) {}
  VariableValue(const MooseVariableBase & var,
                const TagName & tag = Moose::SOLUTION_TAG,
                bool dof = false)
    : _var(var, tag), _dof(dof)
  {
  }

  KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }

  KOKKOS_FUNCTION Real operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const;

private:
  Variable _var;
  bool _dof = false;
};

class VariableGradient
{
public:
  VariableGradient() = default;
  VariableGradient(Variable var) : _var(var) {}
  VariableGradient(const MooseVariableBase & var, const TagName & tag = Moose::SOLUTION_TAG)
    : _var(var, tag)
  {
  }

  KOKKOS_FUNCTION operator bool() const { return _var.coupled(); }

  KOKKOS_FUNCTION Real3 operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const;

private:
  Variable _var;
};

KOKKOS_FUNCTION inline Real
VariableValue::operator()(Datum & datum, unsigned int idx, unsigned int comp) const
{
  KOKKOS_ASSERT(_var.initialized());

  if (_var.coupled())
  {
    auto & sys = datum.system(_var.sys(comp));
    auto var = _var.var(comp);
    auto tag = _var.tag();

    if (_dof)
    {
      unsigned int dof;

      if (datum.isNodal())
      {
        auto node = datum.node();
        dof = sys.getNodeLocalDofIndex(node, 0, var);
      }
      else
      {
        auto elem = datum.elem().id;
        dof = sys.getElemLocalDofIndex(elem, idx, var);
      }

      return sys.getVectorDofValue(dof, tag);
    }
    else
    {
      auto & elem = datum.elem();
      auto side = datum.side();
      auto offset = datum.qpOffset();

      return side == libMesh::invalid_uint ? sys.getVectorQpValue(elem, offset + idx, var, tag)
                                           : sys.getVectorQpValueFace(elem, side, idx, var, tag);
    }
  }
  else
    return _var.value(comp);
}

KOKKOS_FUNCTION inline Real3
VariableGradient::operator()(Datum & datum, unsigned int qp, unsigned int comp) const
{
  KOKKOS_ASSERT(_var.initialized());

  if (_var.coupled())
  {
    KOKKOS_ASSERT(!datum.isNodal());

    auto & elem = datum.elem();
    auto side = datum.side();
    auto offset = datum.qpOffset();

    return side == libMesh::invalid_uint
               ? datum.system(_var.sys(comp))
                     .getVectorQpGrad(elem, offset + qp, _var.var(comp), _var.tag())
               : datum.system(_var.sys(comp))
                     .getVectorQpGradFace(elem, side, datum.J(qp), qp, _var.var(comp), _var.tag());
  }
  else
    return Real3(0);
}

} // namespace Kokkos
} // namespace Moose