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 "MooseError.h"
#include "MooseVariableFieldBase.h"
#include "SystemBase.h"

namespace Moose::Kokkos
{

inline void
checkVariable(const Variable & var, bool expect_vector, const std::string & wrapper_name)
{
  if (!var.initialized())
    mooseError("Attempted to construct Kokkos ", wrapper_name, " with an uninitialized variable.");

  if (var.vector() != expect_vector)
    mooseError("Kokkos ",
               wrapper_name,
               " cannot be constructed with ",
               var.vector() ? "vector" : "scalar",
               " variables.");
}

template <bool is_test>
class VariableShapeValue
{
public:
  KOKKOS_FUNCTION Real operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = is_test ? datum.ife() : 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);
  }
};

template <bool is_test>
class VariableShapeGradient
{
public:
  KOKKOS_FUNCTION const Real3 &
  reference(AssemblyDatum & datum, unsigned int i, unsigned int qp) const;

  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    return datum.J(qp) * reference(datum, i, qp);
  }
};

template <bool is_test>
KOKKOS_FUNCTION const Real3 &
VariableShapeGradient<is_test>::reference(AssemblyDatum & datum,
                                          unsigned int i,
                                          unsigned int qp) const
{
  auto & elem = datum.elem();
  auto side = datum.side();
  auto fe = is_test ? datum.ife() : datum.jfe();

  return 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);
}

using VariablePhiValue = VariableShapeValue<false>;
using VariablePhiGradient = VariableShapeGradient<false>;
using VariableTestValue = VariableShapeValue<true>;
using VariableTestGradient = VariableShapeGradient<true>;
using ADVariablePhiValue = VariablePhiValue;
using ADVariablePhiGradient = VariablePhiGradient;
using ADVariableTestValue = VariableTestValue;
using ADVariableTestGradient = VariableTestGradient;

template <bool is_test>
class VectorVariableShapeValue
{
public:
  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = is_test ? datum.ife() : datum.jfe();

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

template <bool is_test>
class VectorVariableShapeGradient
{
public:
  KOKKOS_FUNCTION const Real33 &
  reference(AssemblyDatum & datum, unsigned int i, unsigned int qp) const;

  KOKKOS_FUNCTION Real33 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    return reference(datum, i, qp) * datum.J(qp).transpose();
  }
};

template <bool is_test>
KOKKOS_FUNCTION const Real33 &
VectorVariableShapeGradient<is_test>::reference(AssemblyDatum & datum,
                                                unsigned int i,
                                                unsigned int qp) const
{
  auto & elem = datum.elem();
  auto side = datum.side();
  auto fe = is_test ? datum.ife() : datum.jfe();

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

template <bool is_test>
class VectorVariableShapeCurl
{
public:
  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum, unsigned int i, unsigned int qp) const
  {
    auto & elem = datum.elem();
    auto side = datum.side();
    auto fe = is_test ? datum.ife() : datum.jfe();

    auto grad =
        side == libMesh::invalid_uint
            ? datum.assembly().getVectorGradPhi(elem.subdomain, elem.type, fe)(i, qp)
            : datum.assembly().getVectorGradPhiFace(elem.subdomain, elem.type, fe)(side)(i, qp);

    return curlFromVectorGradient(grad * datum.J(qp).transpose(), datum.assembly().getDimension());
  }
};

using VectorVariablePhiValue = VectorVariableShapeValue<false>;
using VectorVariablePhiGradient = VectorVariableShapeGradient<false>;
using VectorVariablePhiCurl = VectorVariableShapeCurl<false>;
using VectorVariableTestValue = VectorVariableShapeValue<true>;
using VectorVariableTestGradient = VectorVariableShapeGradient<true>;
using VectorVariableTestCurl = VectorVariableShapeCurl<true>;

template <bool is_ad>
class VariableValueTempl
{
  using real_type = std::conditional_t<is_ad, ADReal, Real>;

public:
  VariableValueTempl() = default;
  VariableValueTempl(Variable var, bool dof = false) : _var(var), _dof(dof)
  {
    checkVariable(_var, false, is_ad ? "ADVariableValue" : "VariableValue");
  }
  VariableValueTempl(const MooseVariableFieldBase & var,
                     const TagName & tag = Moose::SOLUTION_TAG,
                     bool dof = false)
    : _var(var, tag), _dof(dof)
  {
    checkVariable(_var, false, is_ad ? "ADVariableValue" : "VariableValue");
  }
  VariableValueTempl(const std::vector<const MooseVariableFieldBase *> & vars,
                     const TagName & tag = Moose::SOLUTION_TAG,
                     bool dof = false)
    : _var(vars, tag), _dof(dof)
  {
    checkVariable(_var, false, is_ad ? "ADVariableValue" : "VariableValue");
  }
  VariableValueTempl(const std::vector<MooseVariableFieldBase *> & vars,
                     const TagName & tag = Moose::SOLUTION_TAG,
                     bool dof = false)
    : _var(vars, tag), _dof(dof)
  {
    checkVariable(_var, false, is_ad ? "ADVariableValue" : "VariableValue");
  }

  VariableValueTempl(const VariableValueTempl<is_ad> & object);
  VariableValueTempl<is_ad> & operator=(const VariableValueTempl<is_ad> & object);

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

  KOKKOS_FUNCTION auto operator()(Datum & datum, unsigned int idx, unsigned int comp = 0) const
  {
    return get(datum, idx, comp);
  }

  KOKKOS_FUNCTION auto
  operator()(AssemblyDatum & datum, unsigned int idx, unsigned int comp = 0) const;

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

private:
  KOKKOS_FUNCTION auto
  get(Datum & datum, unsigned int idx, unsigned int comp = 0, Real seed = 0) const;

  Variable _var;
  Array<Real> _seed;
  bool _dof = false;
};

template <bool is_ad>
VariableValueTempl<is_ad>::VariableValueTempl(const VariableValueTempl<is_ad> & object)
  : _var(object._var), _seed(object._seed), _dof(object._dof)
{
  if constexpr (is_ad)
    if (_var.coupled())
    {
      if (!_seed.isAlloc())
        _seed.create(_var.components());

      for (unsigned int comp = 0; comp < _var.components(); ++comp)
        _seed[comp] =
            _var.dot() ? _var.mooseVar(comp)->sys().duDotDu(_var.var(comp)) : (_var.old() ? 0 : 1);

      _seed.copyToDevice();
    }
}

template <bool is_ad>
VariableValueTempl<is_ad> &
VariableValueTempl<is_ad>::operator=(const VariableValueTempl<is_ad> & object)
{
  _var = object._var;
  _dof = object._dof;

  return *this;
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableValueTempl<is_ad>::operator()(AssemblyDatum & datum,
                                      unsigned int idx,
                                      unsigned int comp) const
{
  if constexpr (is_ad)
  {
    Real seed =
        datum.do_derivatives() && _var.coupled() && _var.sys(comp) == datum.sys() ? _seed[comp] : 0;

    return get(datum, idx, comp, seed);
  }
  else
    return get(datum, idx, comp);
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableValueTempl<is_ad>::get(Datum & datum,
                               unsigned int idx,
                               unsigned int comp,
                               [[maybe_unused]] Real seed) const
{
  KOKKOS_ASSERT(_var.initialized());

  real_type value;

  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);
      }

      if constexpr (is_ad)
        value = sys.getVectorDofADValue(dof, tag, seed);
      else
        value = sys.getVectorDofValue(dof, tag);
    }
    else
    {
      auto & elem = datum.elem();
      auto side = datum.side();

      if constexpr (is_ad)
        value = side == libMesh::invalid_uint
                    ? sys.getVectorQpADValue(elem, datum.qpOffset(), idx, var, tag, seed)
                    : sys.getVectorQpADValueFace(elem, side, idx, var, tag, seed);
      else
        value = side == libMesh::invalid_uint
                    ? sys.getVectorQpValue(elem, datum.qpOffset() + idx, var, tag)
                    : sys.getVectorQpValueFace(elem, side, idx, var, tag);
    }
  }
  else
    value = _var.value(comp);

  return value;
}

template <bool is_ad>
class VariableGradientTempl
{
  using real3_type = std::conditional_t<is_ad, ADReal3, Real3>;

public:
  VariableGradientTempl() = default;
  VariableGradientTempl(Variable var) : _var(var)
  {
    checkVariable(_var, false, is_ad ? "ADVariableGradient" : "VariableGradient");
  }
  VariableGradientTempl(const MooseVariableFieldBase & var,
                        const TagName & tag = Moose::SOLUTION_TAG)
    : _var(var, tag)
  {
    checkVariable(_var, false, is_ad ? "ADVariableGradient" : "VariableGradient");
  }
  VariableGradientTempl(const std::vector<const MooseVariableFieldBase *> & vars,
                        const TagName & tag = Moose::SOLUTION_TAG)
    : _var(vars, tag)
  {
    checkVariable(_var, false, is_ad ? "ADVariableGradient" : "VariableGradient");
  }
  VariableGradientTempl(const std::vector<MooseVariableFieldBase *> & vars,
                        const TagName & tag = Moose::SOLUTION_TAG)
    : _var(vars, tag)
  {
    checkVariable(_var, false, is_ad ? "ADVariableGradient" : "VariableGradient");
  }

  VariableGradientTempl(const VariableGradientTempl<is_ad> & object);
  VariableGradientTempl<is_ad> & operator=(const VariableGradientTempl<is_ad> & object);

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

  KOKKOS_FUNCTION auto operator()(Datum & datum, unsigned int qp, unsigned int comp = 0) const
  {
    return get(datum, qp, comp);
  }

  KOKKOS_FUNCTION auto
  operator()(AssemblyDatum & datum, unsigned int qp, unsigned int comp = 0) const;

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

private:
  KOKKOS_FUNCTION auto
  get(Datum & datum, unsigned int qp, unsigned int comp = 0, Real seed = 0) const;

  Variable _var;
  Array<Real> _seed;
};

template <bool is_ad>
VariableGradientTempl<is_ad>::VariableGradientTempl(const VariableGradientTempl<is_ad> & object)
  : _var(object._var), _seed(object._seed)
{
  if constexpr (is_ad)
    if (_var.coupled())
    {
      if (!_seed.isAlloc())
        _seed.create(_var.components());

      for (unsigned int comp = 0; comp < _var.components(); ++comp)
        _seed[comp] =
            _var.dot() ? _var.mooseVar(comp)->sys().duDotDu(_var.var(comp)) : (_var.old() ? 0 : 1);

      _seed.copyToDevice();
    }
}

template <bool is_ad>
VariableGradientTempl<is_ad> &
VariableGradientTempl<is_ad>::operator=(const VariableGradientTempl<is_ad> & object)
{
  _var = object._var;

  return *this;
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableGradientTempl<is_ad>::operator()(AssemblyDatum & datum,
                                         unsigned int qp,
                                         unsigned int comp) const
{
  if constexpr (is_ad)
  {
    Real seed =
        datum.do_derivatives() && _var.coupled() && _var.sys(comp) == datum.sys() ? _seed[comp] : 0;

    return get(datum, qp, comp, seed);
  }
  else
    return get(datum, qp, comp);
}

template <bool is_ad>
KOKKOS_FUNCTION auto
VariableGradientTempl<is_ad>::get(Datum & datum,
                                  unsigned int qp,
                                  unsigned int comp,
                                  [[maybe_unused]] Real seed) const
{
  KOKKOS_ASSERT(_var.initialized());

  real3_type grad;

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

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

    if constexpr (is_ad)
      grad =
          side == libMesh::invalid_uint
              ? datum.system(_var.sys(comp))
                    .getVectorQpADGrad(
                        elem, datum.J(qp), datum.qpOffset(), qp, _var.var(comp), _var.tag(), seed)
              : datum.system(_var.sys(comp))
                    .getVectorQpADGradFace(
                        elem, side, datum.J(qp), qp, _var.var(comp), _var.tag(), seed);
    else
      grad =
          side == libMesh::invalid_uint
              ? datum.system(_var.sys(comp))
                    .getVectorQpGrad(elem, datum.qpOffset() + qp, _var.var(comp), _var.tag())
              : datum.system(_var.sys(comp))
                    .getVectorQpGradFace(elem, side, datum.J(qp), qp, _var.var(comp), _var.tag());
  }

  return grad;
}

using VariableValue = VariableValueTempl<false>;
using ADVariableValue = VariableValueTempl<true>;
using VariableGradient = VariableGradientTempl<false>;
using ADVariableGradient = VariableGradientTempl<true>;

class VectorVariableValue
{
public:
  VectorVariableValue() = default;
  VectorVariableValue(Variable var, bool dof = false) : _var(var), _dof(dof)
  {
    checkVariable(_var, true, "VectorVariableValue");
  }
  VectorVariableValue(const MooseVariableFieldBase & var,
                      const TagName & tag = Moose::SOLUTION_TAG,
                      bool dof = false)
    : _var(var, tag), _dof(dof)
  {
    checkVariable(_var, true, "VectorVariableValue");
  }

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

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

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

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

class VectorVariableGradient
{
public:
  VectorVariableGradient() = default;
  VectorVariableGradient(Variable var) : _var(var)
  {
    checkVariable(_var, true, "VectorVariableGradient");
  }
  VectorVariableGradient(const MooseVariableFieldBase & var,
                         const TagName & tag = Moose::SOLUTION_TAG)
    : _var(var, tag)
  {
    checkVariable(_var, true, "VectorVariableGradient");
  }

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

  KOKKOS_FUNCTION Real33 operator()(AssemblyDatum & datum,
                                    unsigned int qp,
                                    unsigned int comp = 0) const;

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

private:
  Variable _var;
};

class VectorVariableCurl
{
public:
  VectorVariableCurl() = default;
  VectorVariableCurl(Variable var) : _var(var) { checkVariable(_var, true, "VectorVariableCurl"); }
  VectorVariableCurl(const MooseVariableFieldBase & var, const TagName & tag = Moose::SOLUTION_TAG)
    : _var(var, tag)
  {
    checkVariable(_var, true, "VectorVariableCurl");
  }

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

  KOKKOS_FUNCTION Real3 operator()(AssemblyDatum & datum,
                                   unsigned int qp,
                                   unsigned int comp = 0) const;

  KOKKOS_FUNCTION const Variable & variable() const { return _var; }

private:
  Variable _var;
};

KOKKOS_FUNCTION inline Real3
VectorVariableValue::operator()(AssemblyDatum & datum, unsigned int idx, unsigned int comp) const
{
  KOKKOS_ASSERT(_var.initialized());

  Real3 value = 0;

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

    if (_dof)
    {
      auto dimension = datum.assembly().getDimension();

      if (datum.isNodal())
      {
        auto node = datum.node();

        for (unsigned int c = 0; c < dimension; ++c)
          value(c) = sys.getVectorDofValue(sys.getNodeLocalDofIndex(node, c, var), tag);
      }
      else
      {
        auto elem = datum.elem().id;
        auto offset = idx * dimension;

        for (unsigned int c = 0; c < dimension; ++c)
          value(c) = sys.getVectorDofValue(sys.getElemLocalDofIndex(elem, offset + c, var), tag);
      }
    }
    else
    {
      KOKKOS_ASSERT(!datum.isNodal());

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

      if (side == libMesh::invalid_uint)
        value = sys.getVectorQpVectorValue(elem, datum.qpOffset() + idx, var, tag);
      else
        value = sys.getVectorQpVectorValueFace(elem, side, idx, var, tag);
    }
  }
  else
    value = _var.vectorValue(comp);

  return value;
}

KOKKOS_FUNCTION inline Real33
VectorVariableGradient::operator()(AssemblyDatum & datum, unsigned int qp, unsigned int comp) const
{
  KOKKOS_ASSERT(_var.initialized());

  Real33 grad = 0;

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

    auto & elem = datum.elem();
    auto side = datum.side();
    auto & sys = datum.system(_var.sys(comp));
    auto var = _var.var(comp);
    auto tag = _var.tag();

    if (side == libMesh::invalid_uint)
      grad = sys.getVectorQpVectorGrad(elem, datum.qpOffset() + qp, var, tag);
    else
      grad = sys.getVectorQpVectorGradFace(elem, side, datum.J(qp), qp, var, tag);
  }

  return grad;
}

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

  Real3 curl = 0;

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

    auto & elem = datum.elem();
    auto side = datum.side();
    auto & sys = datum.system(_var.sys(comp));
    auto var = _var.var(comp);
    auto tag = _var.tag();

    if (side == libMesh::invalid_uint)
      curl = sys.getVectorQpVectorCurl(elem, datum.qpOffset() + qp, var, tag);
    else
    {
      auto fe = sys.getFETypeID(var);
      auto n_dofs = datum.assembly().getNumDofs(elem.type, fe);
      auto & grad_phi = datum.assembly().getVectorGradPhiFace(elem.subdomain, elem.type, fe)(side);
      auto jacobian = datum.J(qp);
      auto jacobian_transpose = jacobian.transpose();
      Real33 grad = 0;

      for (unsigned int i = 0; i < n_dofs; ++i)
        grad += sys.getVectorDofValue(sys.getElemLocalDofIndex(elem.id, i, var), tag) *
                (grad_phi(i, qp) * jacobian_transpose);

      curl = curlFromVectorGradient(grad, datum.assembly().getDimension());
    }
  }

  return curl;
}

template <>
struct ArrayDeepCopy<ADVariableValue>
{
  static constexpr bool value = true;
};

template <>
struct ArrayDeepCopy<ADVariableGradient>
{
  static constexpr bool value = true;
};

} // namespace Moose::Kokkos