functional.hpp

Go to the documentation of this file.

// Copyright (c) 2019-2024, Lawrence Livermore National Security, LLC and
// other Serac Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include "mfem.hpp"
 
#include "serac/infrastructure/logger.hpp"
#include "serac/numerics/functional/tensor.hpp"
#include "serac/numerics/functional/quadrature.hpp"
#include "serac/numerics/functional/finite_element.hpp"
#include "serac/numerics/functional/integral.hpp"
#include "serac/numerics/functional/dof_numbering.hpp"
#include "serac/numerics/functional/differentiate_wrt.hpp"
 
#include "serac/numerics/functional/element_restriction.hpp"
 
#include "serac/numerics/functional/domain.hpp"
 
#include <array>
#include <vector>
 
namespace serac {
 
template <int... i>
struct DependsOn {
};
 
template <typename... T>
constexpr uint32_t index_of_differentiation()
{
  constexpr uint32_t n          = sizeof...(T);
  bool               matching[] = {std::is_same_v<T, differentiate_wrt_this>...};
  for (uint32_t i = 0; i < n; i++) {
    if (matching[i]) {
      return i;
    }
  }
  return NO_DIFFERENTIATION;
}
 
template <int ind>
struct Index {
  constexpr operator int() { return ind; }
};
 
inline void check_for_missing_nodal_gridfunc(const mfem::Mesh& mesh)
{
  if (mesh.GetNodes() == nullptr) {
    SLIC_ERROR_ROOT(
        R"errmsg(
      the provided mesh does not have a nodal gridfunction. 
      If you created an mfem::Mesh manually, make sure that the 
      following member functions are invoked before use
 
      > mfem::Mesh::EnsureNodes();
      > mfem::Mesh::ExchangeFaceNbrData();
 
      or else the mfem::Mesh won't be fully initialized
      )errmsg";);
  }
}
 
inline void check_for_unsupported_elements(const mfem::Mesh& mesh)
{
  int num_elements = mesh.GetNE();
  for (int e = 0; e < num_elements; e++) {
    auto type = mesh.GetElementType(e);
    if (type == mfem::Element::POINT || type == mfem::Element::WEDGE || type == mfem::Element::PYRAMID) {
      SLIC_ERROR_ROOT("Mesh contains unsupported element type");
    }
  }
}
 
template <typename function_space>
inline std::pair<std::unique_ptr<mfem::ParFiniteElementSpace>, std::unique_ptr<mfem::FiniteElementCollection>>
generateParFiniteElementSpace(mfem::ParMesh* mesh)
{
  const int                                      dim = mesh->Dimension();
  std::unique_ptr<mfem::FiniteElementCollection> fec;
  const auto                                     ordering = mfem::Ordering::byNODES;
 
  switch (function_space::family) {
    case Family::H1:
      fec = std::make_unique<mfem::H1_FECollection>(function_space::order, dim);
      break;
    case Family::HCURL:
      fec = std::make_unique<mfem::ND_FECollection>(function_space::order, dim);
      break;
    case Family::HDIV:
      fec = std::make_unique<mfem::RT_FECollection>(function_space::order, dim);
      break;
    case Family::L2:
      // We use GaussLobatto basis functions as this is what is used for the serac::Functional FE kernels
      fec = std::make_unique<mfem::L2_FECollection>(function_space::order, dim, mfem::BasisType::GaussLobatto);
      break;
    default:
      return std::pair<std::unique_ptr<mfem::ParFiniteElementSpace>, std::unique_ptr<mfem::FiniteElementCollection>>(
          nullptr, nullptr);
      break;
  }
 
  auto fes = std::make_unique<mfem::ParFiniteElementSpace>(mesh, fec.get(), function_space::components, ordering);
 
  return std::pair(std::move(fes), std::move(fec));
}
 
template <typename T, ExecutionSpace exec = serac::default_execution_space>
class Functional;
 
template <typename test, typename... trials, ExecutionSpace exec>
class Functional<test(trials...), exec> {
  static constexpr tuple<trials...> trial_spaces{};
  static constexpr uint32_t         num_trial_spaces = sizeof...(trials);
  static constexpr auto             Q                = std::max({test::order, trials::order...}) + 1;
 
  static constexpr mfem::Geometry::Type elem_geom[4]    = {mfem::Geometry::INVALID, mfem::Geometry::SEGMENT,
                                                           mfem::Geometry::SQUARE, mfem::Geometry::CUBE};
  static constexpr mfem::Geometry::Type simplex_geom[4] = {mfem::Geometry::INVALID, mfem::Geometry::SEGMENT,
                                                           mfem::Geometry::TRIANGLE, mfem::Geometry::TETRAHEDRON};
 
  class Gradient;
 
  // clang-format off
  template <uint32_t i> 
  struct operator_paren_return {
    using type = typename std::conditional<
        i == NO_DIFFERENTIATION,               // if `i` indicates that we want to skip differentiation
        mfem::Vector&,                         // we just return the value
        serac::tuple<mfem::Vector&, Gradient&> // otherwise we return the value and the derivative w.r.t arg `i`
        >::type;
  };
  // clang-format on
 
public:
  Functional(const mfem::ParFiniteElementSpace*                               test_fes,
             std::array<const mfem::ParFiniteElementSpace*, num_trial_spaces> trial_fes)
      : update_qdata_(false), test_space_(test_fes), trial_space_(trial_fes)
  {
    auto mem_type = mfem::Device::GetMemoryType();
 
    for (auto type : {Domain::Type::Elements, Domain::Type::BoundaryElements}) {
      input_E_[type].resize(num_trial_spaces);
    }
 
    for (uint32_t i = 0; i < num_trial_spaces; i++) {
      P_trial_[i] = trial_space_[i]->GetProlongationMatrix();
 
      input_L_[i].SetSize(P_trial_[i]->Height(), mfem::Device::GetMemoryType());
 
      // L->E
      for (auto type : {Domain::Type::Elements, Domain::Type::BoundaryElements}) {
        if (type == Domain::Type::Elements) {
          G_trial_[type][i] = BlockElementRestriction(trial_fes[i]);
        } else {
          G_trial_[type][i] = BlockElementRestriction(trial_fes[i], FaceType::BOUNDARY);
        }
 
        // note: we have to use "Update" here, as mfem::BlockVector's
        // copy assignment ctor (operator=) doesn't let you make changes
        // to the block size
        input_E_[type][i].Update(G_trial_[type][i].bOffsets(), mem_type);
      }
    }
 
    for (auto type : {Domain::Type::Elements, Domain::Type::BoundaryElements}) {
      if (type == Domain::Type::Elements) {
        G_test_[type] = BlockElementRestriction(test_fes);
      } else {
        G_test_[type] = BlockElementRestriction(test_fes, FaceType::BOUNDARY);
      }
 
      output_E_[type].Update(G_test_[type].bOffsets(), mem_type);
    }
 
    P_test_ = test_space_->GetProlongationMatrix();
 
    output_L_.SetSize(P_test_->Height(), mem_type);
 
    output_T_.SetSize(test_fes->GetTrueVSize(), mem_type);
 
    // gradient objects depend on some member variables in
    // Functional, so we initialize the gradient objects last
    // to ensure that those member variables are initialized first
    for (uint32_t i = 0; i < num_trial_spaces; i++) {
      grad_.emplace_back(*this, i);
    }
  }
 
  template <int dim, int... args, typename lambda, typename qpt_data_type = Nothing>
  void AddDomainIntegral(Dimension<dim>, DependsOn<args...>, lambda&& integrand, mfem::Mesh& domain,
                         std::shared_ptr<QuadratureData<qpt_data_type>> qdata = NoQData)
  {
    if (domain.GetNE() == 0) return;
 
    SLIC_ERROR_ROOT_IF(dim != domain.Dimension(), "invalid mesh dimension for domain integral");
 
    check_for_unsupported_elements(domain);
    check_for_missing_nodal_gridfunc(domain);
 
    using signature = test(decltype(serac::type<args>(trial_spaces))...);
    integrals_.push_back(
        MakeDomainIntegral<signature, Q, dim>(EntireDomain(domain), integrand, qdata, std::vector<uint32_t>{args...}));
  }
 
  template <int dim, int... args, typename lambda, typename qpt_data_type = Nothing>
  void AddDomainIntegral(Dimension<dim>, DependsOn<args...>, lambda&& integrand, Domain& domain,
                         std::shared_ptr<QuadratureData<qpt_data_type>> qdata = NoQData)
  {
    if (domain.mesh_.GetNE() == 0) return;
 
    SLIC_ERROR_ROOT_IF(dim != domain.mesh_.Dimension(), "invalid mesh dimension for domain integral");
 
    check_for_unsupported_elements(domain.mesh_);
    check_for_missing_nodal_gridfunc(domain.mesh_);
 
    using signature = test(decltype(serac::type<args>(trial_spaces))...);
    integrals_.push_back(
        MakeDomainIntegral<signature, Q, dim>(domain, integrand, qdata, std::vector<uint32_t>{args...}));
  }
 
  template <int dim, int... args, typename lambda>
  void AddBoundaryIntegral(Dimension<dim>, DependsOn<args...>, lambda&& integrand, mfem::Mesh& domain)
  {
    auto num_bdr_elements = domain.GetNBE();
    if (num_bdr_elements == 0) return;
 
    check_for_missing_nodal_gridfunc(domain);
 
    using signature = test(decltype(serac::type<args>(trial_spaces))...);
    integrals_.push_back(
        MakeBoundaryIntegral<signature, Q, dim>(EntireBoundary(domain), integrand, std::vector<uint32_t>{args...}));
  }
 
  template <int dim, int... args, typename lambda>
  void AddBoundaryIntegral(Dimension<dim>, DependsOn<args...>, lambda&& integrand, const Domain& domain)
  {
    auto num_bdr_elements = domain.mesh_.GetNBE();
    if (num_bdr_elements == 0) return;
 
    SLIC_ERROR_ROOT_IF(dim != domain.dim_, "invalid domain of integration for boundary integral");
 
    check_for_missing_nodal_gridfunc(domain.mesh_);
 
    using signature = test(decltype(serac::type<args>(trial_spaces))...);
    integrals_.push_back(MakeBoundaryIntegral<signature, Q, dim>(domain, integrand, std::vector<uint32_t>{args...}));
  }
 
  template <int... args, typename lambda, typename qpt_data_type = Nothing>
  void AddAreaIntegral(DependsOn<args...> which_args, lambda&& integrand, mfem::Mesh& domain,
                       std::shared_ptr<QuadratureData<qpt_data_type>> data = NoQData)
  {
    AddDomainIntegral(Dimension<2>{}, which_args, integrand, domain, data);
  }
 
  template <int... args, typename lambda, typename qpt_data_type = Nothing>
  void AddVolumeIntegral(DependsOn<args...> which_args, lambda&& integrand, mfem::Mesh& domain,
                         std::shared_ptr<QuadratureData<qpt_data_type>> data = NoQData)
  {
    AddDomainIntegral(Dimension<3>{}, which_args, integrand, domain, data);
  }
 
  template <int... args, typename lambda>
  void AddSurfaceIntegral(DependsOn<args...> which_args, lambda&& integrand, mfem::Mesh& domain)
  {
    AddBoundaryIntegral(Dimension<2>{}, which_args, integrand, domain);
  }
 
  void ActionOfGradient(const mfem::Vector& input_T, mfem::Vector& output_T, uint32_t which) const
  {
    P_trial_[which]->Mult(input_T, input_L_[which]);
 
    output_L_ = 0.0;
 
    // this is used to mark when gather operations have been performed,
    // to avoid doing them more than once per trial space
    bool already_computed[Domain::num_types]{};  // default initializes to `false`
 
    for (auto& integral : integrals_) {
      auto type = integral.domain_.type_;
 
      if (!already_computed[type]) {
        G_trial_[type][which].Gather(input_L_[which], input_E_[type][which]);
        already_computed[type] = true;
      }
 
      integral.GradientMult(input_E_[type][which], output_E_[type], which);
 
      // scatter-add to compute residuals on the local processor
      G_test_[type].ScatterAdd(output_E_[type], output_L_);
    }
 
    // scatter-add to compute global residuals
    P_test_->MultTranspose(output_L_, output_T);
  }
 
  template <uint32_t wrt, typename... T>
  typename operator_paren_return<wrt>::type operator()(DifferentiateWRT<wrt>, double t, const T&... args)
  {
    const mfem::Vector* input_T[] = {&static_cast<const mfem::Vector&>(args)...};
 
    // get the values for each local processor
    for (uint32_t i = 0; i < num_trial_spaces; i++) {
      P_trial_[i]->Mult(*input_T[i], input_L_[i]);
    }
 
    output_L_ = 0.0;
 
    // this is used to mark when operations have been performed,
    // to avoid doing them more than once
    bool already_computed[Domain::num_types][num_trial_spaces]{};  // default initializes to `false`
 
    for (auto& integral : integrals_) {
      auto type = integral.domain_.type_;
 
      for (auto i : integral.active_trial_spaces_) {
        if (!already_computed[type][i]) {
          G_trial_[type][i].Gather(input_L_[i], input_E_[type][i]);
          already_computed[type][i] = true;
        }
      }
 
      integral.Mult(t, input_E_[type], output_E_[type], wrt, update_qdata_);
 
      // scatter-add to compute residuals on the local processor
      G_test_[type].ScatterAdd(output_E_[type], output_L_);
    }
 
    // scatter-add to compute global residuals
    P_test_->MultTranspose(output_L_, output_T_);
 
    if constexpr (wrt != NO_DIFFERENTIATION) {
      // if the user has indicated they'd like to evaluate and differentiate w.r.t.
      // a specific argument, then we return both the value and gradient w.r.t. that argument
      //
      // mfem::Vector arg0 = ...;
      // mfem::Vector arg1 = ...;
      // e.g. auto [value, gradient_wrt_arg1] = my_functional(arg0, differentiate_wrt(arg1));
      return {output_T_, grad_[wrt]};
    }
    if constexpr (wrt == NO_DIFFERENTIATION) {
      // if the user passes only `mfem::Vector`s then we assume they only want the output value
      //
      // mfem::Vector arg0 = ...;
      // mfem::Vector arg1 = ...;
      // e.g. mfem::Vector value = my_functional(arg0, arg1);
      return output_T_;
    }
  }
 
  template <typename... T>
  auto operator()(double t, const T&... args)
  {
    constexpr int num_differentiated_arguments = (std::is_same_v<T, differentiate_wrt_this> + ...);
    static_assert(num_differentiated_arguments <= 1,
                  "Error: Functional::operator() can only differentiate w.r.t. 1 argument a time");
    static_assert(sizeof...(T) == num_trial_spaces,
                  "Error: Functional::operator() must take exactly as many arguments as trial spaces");
 
    [[maybe_unused]] constexpr uint32_t i = index_of_differentiation<T...>();
 
    return (*this)(DifferentiateWRT<i>{}, t, args...);
  }
 
  void updateQdata(bool update_flag) { update_qdata_ = update_flag; }
 
private:
  bool update_qdata_;
 
  class Gradient : public mfem::Operator {
  public:
    Gradient(Functional<test(trials...), exec>& f, uint32_t which = 0)
        : mfem::Operator(f.test_space_->GetTrueVSize(), f.trial_space_[which]->GetTrueVSize()),
          form_(f),
          lookup_tables(f.G_test_[Domain::Type::Elements], f.G_trial_[Domain::Type::Elements][which]),
          which_argument(which),
          test_space_(f.test_space_),
          trial_space_(f.trial_space_[which]),
          df_(f.test_space_->GetTrueVSize())
    {
    }
 
    virtual void Mult(const mfem::Vector& dx, mfem::Vector& df) const override
    {
      form_.ActionOfGradient(dx, df, which_argument);
    }
 
    mfem::Vector& operator()(const mfem::Vector& dx)
    {
      form_.ActionOfGradient(dx, df_, which_argument);
      return df_;
    }
 
    std::unique_ptr<mfem::HypreParMatrix> assemble()
    {
      // the CSR graph (sparsity pattern) is reusable, so we cache
      // that and ask mfem to not free that memory in ~SparseMatrix()
      constexpr bool sparse_matrix_frees_graph_ptrs = false;
 
      // the CSR values are NOT reusable, so we pass ownership of
      // them to the mfem::SparseMatrix, to be freed in ~SparseMatrix()
      constexpr bool sparse_matrix_frees_values_ptr = true;
 
      constexpr bool col_ind_is_sorted = true;
 
      double* values = new double[lookup_tables.nnz]{};
 
      std::map<mfem::Geometry::Type, ExecArray<double, 3, exec>> element_gradients[Domain::num_types];
 
      for (auto& integral : form_.integrals_) {
        auto& K_elem             = element_gradients[integral.domain_.type_];
        auto& test_restrictions  = form_.G_test_[integral.domain_.type_].restrictions;
        auto& trial_restrictions = form_.G_trial_[integral.domain_.type_][which_argument].restrictions;
 
        if (K_elem.empty()) {
          for (auto& [geom, test_restriction] : test_restrictions) {
            auto& trial_restriction = trial_restrictions[geom];
 
            K_elem[geom] = ExecArray<double, 3, exec>(test_restriction.num_elements,
                                                      trial_restriction.nodes_per_elem * trial_restriction.components,
                                                      test_restriction.nodes_per_elem * test_restriction.components);
 
            detail::zero_out(K_elem[geom]);
          }
        }
 
        integral.ComputeElementGradients(K_elem, which_argument);
      }
 
      for (auto type : {Domain::Type::Elements, Domain::Type::BoundaryElements}) {
        auto& K_elem             = element_gradients[type];
        auto& test_restrictions  = form_.G_test_[type].restrictions;
        auto& trial_restrictions = form_.G_trial_[type][which_argument].restrictions;
 
        if (!K_elem.empty()) {
          for (auto [geom, elem_matrices] : K_elem) {
            std::vector<DoF> test_vdofs(test_restrictions[geom].nodes_per_elem * test_restrictions[geom].components);
            std::vector<DoF> trial_vdofs(trial_restrictions[geom].nodes_per_elem * trial_restrictions[geom].components);
 
            for (axom::IndexType e = 0; e < elem_matrices.shape()[0]; e++) {
              test_restrictions[geom].GetElementVDofs(e, test_vdofs);
              trial_restrictions[geom].GetElementVDofs(e, trial_vdofs);
 
              for (uint32_t i = 0; i < uint32_t(elem_matrices.shape()[1]); i++) {
                int col = int(trial_vdofs[i].index());
 
                for (uint32_t j = 0; j < uint32_t(elem_matrices.shape()[2]); j++) {
                  int row = int(test_vdofs[j].index());
 
                  int sign = test_vdofs[j].sign() * trial_vdofs[i].sign();
 
                  // note: col / row appear backwards here, because the element matrix kernel
                  //       is actually transposed, as a result of being row-major storage.
                  //
                  //       This is kind of confusing, and will be fixed in a future refactor
                  //       of the element gradient kernel implementation
                  [[maybe_unused]] auto nz = lookup_tables(row, col);
                  values[lookup_tables(row, col)] += sign * elem_matrices(e, i, j);
                }
              }
            }
          }
        }
      }
 
      // Copy the column indices to an auxilliary array as MFEM can mutate these during HypreParMatrix construction
      col_ind_copy_ = lookup_tables.col_ind;
 
      auto J_local =
          mfem::SparseMatrix(lookup_tables.row_ptr.data(), col_ind_copy_.data(), values, form_.output_L_.Size(),
                             form_.input_L_[which_argument].Size(), sparse_matrix_frees_graph_ptrs,
                             sparse_matrix_frees_values_ptr, col_ind_is_sorted);
 
      auto* R = form_.test_space_->Dof_TrueDof_Matrix();
 
      auto* A =
          new mfem::HypreParMatrix(test_space_->GetComm(), test_space_->GlobalVSize(), trial_space_->GlobalVSize(),
                                   test_space_->GetDofOffsets(), trial_space_->GetDofOffsets(), &J_local);
 
      auto* P = trial_space_->Dof_TrueDof_Matrix();
 
      std::unique_ptr<mfem::HypreParMatrix> K(mfem::RAP(R, A, P));
 
      delete A;
 
      return K;
    };
 
    friend auto assemble(Gradient& g) { return g.assemble(); }
 
  private:
    Functional<test(trials...), exec>& form_;
 
    GradientAssemblyLookupTables lookup_tables;
 
    std::vector<int> col_ind_copy_;
 
    uint32_t which_argument;
 
    const mfem::ParFiniteElementSpace* test_space_;
 
    const mfem::ParFiniteElementSpace* trial_space_;
 
    mfem::Vector df_;
  };
 
  const mfem::ParFiniteElementSpace* test_space_;
 
  std::array<const mfem::ParFiniteElementSpace*, num_trial_spaces> trial_space_;
 
  const mfem::Operator* P_trial_[num_trial_spaces];
 
  mutable mfem::Vector input_L_[num_trial_spaces];
 
  BlockElementRestriction G_trial_[Domain::num_types][num_trial_spaces];
 
  mutable std::vector<mfem::BlockVector> input_E_[Domain::num_types];
 
  std::vector<Integral> integrals_;
 
  mutable mfem::BlockVector output_E_[Domain::num_types];
 
  BlockElementRestriction G_test_[Domain::num_types];
 
  mutable mfem::Vector output_L_;
 
  const mfem::Operator* P_test_;
 
  mutable mfem::Vector output_T_;
 
  mutable std::vector<Gradient> grad_;
};
 
}  // namespace serac
 
#include "functional_qoi.inl"