integral.hpp

// Copyright (c) Lawrence Livermore National Security, LLC and
// other Smith Project Developers. See the top-level LICENSE file for
// details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
 
#pragma once
 
#include <array>
#include <memory>
 
#include "mfem.hpp"
 
#include "smith/infrastructure/accelerator.hpp"
#include "smith/numerics/functional/geometric_factors.hpp"
#include "smith/numerics/functional/function_signature.hpp"
#include "smith/numerics/functional/domain_integral_kernels.hpp"
#include "smith/numerics/functional/boundary_integral_kernels.hpp"
#include "smith/numerics/functional/interior_face_integral_kernels.hpp"
#include "smith/numerics/functional/differentiate_wrt.hpp"
 
namespace smith {
 
struct Integral {
  static constexpr std::size_t num_types = 3;
 
  Integral(const Domain& d, std::vector<uint32_t> trial_space_indices)
      : domain_(d), active_trial_spaces_(trial_space_indices)
  {
    std::size_t num_trial_spaces = trial_space_indices.size();
    evaluation_with_AD_.resize(num_trial_spaces);
    jvp_.resize(num_trial_spaces);
    element_gradient_.resize(num_trial_spaces);
 
    for (uint32_t i = 0; i < num_trial_spaces; i++) {
      functional_to_integral_index_[active_trial_spaces_[i]] = i;
    }
  }
 
  void Mult(double t, const std::vector<mfem::BlockVector>& input_E, mfem::BlockVector& output_E,
            uint32_t differentiation_index, bool update_state) const
  {
    output_E = 0.0;
 
    bool with_AD =
        (functional_to_integral_index_.count(differentiation_index) > 0 && differentiation_index != NO_DIFFERENTIATION);
    auto& kernels =
        (with_AD) ? evaluation_with_AD_[functional_to_integral_index_.at(differentiation_index)] : evaluation_;
    for (auto& [geometry, func] : kernels) {
      std::vector<const double*> inputs(active_trial_spaces_.size());
      for (std::size_t i = 0; i < active_trial_spaces_.size(); i++) {
        inputs[i] = input_E[uint32_t(active_trial_spaces_[i])].GetBlock(geometry).Read();
      }
      func(t, inputs, output_E.GetBlock(geometry).ReadWrite(), update_state);
    }
  }
 
  void GradientMult(const mfem::BlockVector& dinput_E, mfem::BlockVector& doutput_E,
                    uint32_t differentiation_index) const
  {
    doutput_E = 0.0;
 
    // if this integral actually depends on the specified variable
    if (functional_to_integral_index_.count(differentiation_index) > 0) {
      for (auto& [geometry, func] : jvp_[functional_to_integral_index_.at(differentiation_index)]) {
        func(dinput_E.GetBlock(geometry).Read(), doutput_E.GetBlock(geometry).ReadWrite());
      }
    }
  }
 
  void ComputeElementGradients(std::map<mfem::Geometry::Type, ExecArray<double, 3, ExecutionSpace::CPU> >& K_e,
                               uint32_t differentiation_index) const
  {
    // if this integral actually depends on the specified variable
    if (functional_to_integral_index_.count(differentiation_index) > 0) {
      for (auto& [geometry, func] : element_gradient_[functional_to_integral_index_.at(differentiation_index)]) {
        func(view(K_e[geometry]));
      }
    }
  }
 
  bool DependsOn(uint32_t which) const
  {
    for (uint32_t i : active_trial_spaces_) {
      if (which == i) return true;
    }
    return false;
  }
 
  Domain domain_;
 
  using eval_func = std::function<void(double, const std::vector<const double*>&, double*, bool)>;
 
  std::map<mfem::Geometry::Type, eval_func> evaluation_;
 
  std::vector<std::map<mfem::Geometry::Type, eval_func> > evaluation_with_AD_;
 
  using jacobian_vector_product_func = std::function<void(const double*, double*)>;
 
  std::vector<std::map<mfem::Geometry::Type, jacobian_vector_product_func> > jvp_;
 
  using grad_func = std::function<void(ExecArrayView<double, 3, ExecutionSpace::CPU>)>;
 
  std::vector<std::map<mfem::Geometry::Type, grad_func> > element_gradient_;
 
  std::vector<uint32_t> active_trial_spaces_;
 
  std::map<uint32_t, uint32_t> functional_to_integral_index_;
 
  std::map<mfem::Geometry::Type, GeometricFactors> geometric_factors_;
};
 
template <mfem::Geometry::Type geom, int Q, typename test, typename... trials, typename lambda_type,
          typename qpt_data_type>
void generate_kernels(FunctionSignature<test(trials...)> s, Integral& integral, const lambda_type& qf,
                      std::shared_ptr<QuadratureData<qpt_data_type> > qdata)
{
  integral.geometric_factors_[geom] = GeometricFactors(integral.domain_, Q, geom);
  GeometricFactors& gf = integral.geometric_factors_[geom];
  if (gf.num_elements == 0) return;
 
  const double* positions = gf.X.Read();
  const double* jacobians = gf.J.Read();
  const uint32_t num_elements = uint32_t(gf.num_elements);
  const uint32_t qpts_per_element = num_quadrature_points(geom, Q);
 
  std::shared_ptr<zero> dummy_derivatives;
  integral.evaluation_[geom] = domain_integral::evaluation_kernel<NO_DIFFERENTIATION, Q, geom>(
      s, qf, positions, jacobians, qdata, dummy_derivatives, num_elements);
 
  constexpr std::size_t num_args = s.num_args;
  [[maybe_unused]] static constexpr int dim = dimension_of(geom);
  for_constexpr<num_args>([&](auto index) {
    // allocate memory for the derivatives of the q-function at each quadrature point
    //
    // Note: ptrs' lifetime is managed in an unusual way! It is captured by-value in the
    // action_of_gradient functor below to augment the reference count, and extend its lifetime to match
    // that of the DomainIntegral that allocated it.
    using derivative_type = decltype(domain_integral::get_derivative_type<index, dim, trials...>(qf, qpt_data_type{}));
    auto ptr = accelerator::make_shared_array<ExecutionSpace::CPU, derivative_type>(num_elements * qpts_per_element);
 
    integral.evaluation_with_AD_[index][geom] =
        domain_integral::evaluation_kernel<index, Q, geom>(s, qf, positions, jacobians, qdata, ptr, num_elements);
 
    integral.jvp_[index][geom] = domain_integral::jacobian_vector_product_kernel<index, Q, geom>(s, ptr, num_elements);
    integral.element_gradient_[index][geom] =
        domain_integral::element_gradient_kernel<index, Q, geom>(s, ptr, num_elements);
  });
}
 
template <typename s, int Q, int dim, typename lambda_type, typename qpt_data_type>
Integral MakeDomainIntegral(const Domain& domain, const lambda_type& qf,
                            std::shared_ptr<QuadratureData<qpt_data_type> > qdata,
                            std::vector<uint32_t> argument_indices)
{
  FunctionSignature<s> signature;
 
  SLIC_ERROR_IF(domain.type_ != Domain::Type::Elements, "Error: trying to evaluate a domain integral over a boundary");
 
  Integral integral(domain, argument_indices);
 
  if constexpr (dim == 2) {
    generate_kernels<mfem::Geometry::TRIANGLE, Q>(signature, integral, qf, qdata);
    generate_kernels<mfem::Geometry::SQUARE, Q>(signature, integral, qf, qdata);
  }
 
  if constexpr (dim == 3) {
    generate_kernels<mfem::Geometry::TETRAHEDRON, Q>(signature, integral, qf, qdata);
    generate_kernels<mfem::Geometry::CUBE, Q>(signature, integral, qf, qdata);
  }
 
  return integral;
}
 
template <mfem::Geometry::Type geom, int Q, typename test, typename... trials, typename lambda_type>
void generate_bdr_kernels(FunctionSignature<test(trials...)> s, Integral& integral, const lambda_type& qf)
{
  integral.geometric_factors_[geom] = GeometricFactors(integral.domain_, Q, geom);
  GeometricFactors& gf = integral.geometric_factors_[geom];
  if (gf.num_elements == 0) return;
 
  const double* positions = gf.X.Read();
  const double* jacobians = gf.J.Read();
  const uint32_t num_elements = uint32_t(gf.num_elements);
  const uint32_t qpts_per_element = num_quadrature_points(geom, Q);
 
  std::shared_ptr<zero> dummy_derivatives;
  integral.evaluation_[geom] = boundary_integral::evaluation_kernel<NO_DIFFERENTIATION, Q, geom>(
      s, qf, positions, jacobians, dummy_derivatives, num_elements);
 
  constexpr std::size_t num_args = s.num_args;
  [[maybe_unused]] static constexpr int dim = dimension_of(geom);
  for_constexpr<num_args>([&](auto index) {
    // allocate memory for the derivatives of the q-function at each quadrature point
    //
    // Note: ptrs' lifetime is managed in an unusual way! It is captured by-value in the
    // action_of_gradient functor below to augment the reference count, and extend its lifetime to match
    // that of the boundaryIntegral that allocated it.
    using derivative_type = decltype(boundary_integral::get_derivative_type<index, dim, trials...>(qf));
    auto ptr = accelerator::make_shared_array<ExecutionSpace::CPU, derivative_type>(num_elements * qpts_per_element);
 
    integral.evaluation_with_AD_[index][geom] =
        boundary_integral::evaluation_kernel<index, Q, geom>(s, qf, positions, jacobians, ptr, num_elements);
 
    integral.jvp_[index][geom] =
        boundary_integral::jacobian_vector_product_kernel<index, Q, geom>(s, ptr, num_elements);
    integral.element_gradient_[index][geom] =
        boundary_integral::element_gradient_kernel<index, Q, geom>(s, ptr, num_elements);
  });
}
 
template <typename s, int Q, int dim, typename lambda_type>
Integral MakeBoundaryIntegral(const Domain& domain, const lambda_type& qf, std::vector<uint32_t> argument_indices)
{
  FunctionSignature<s> signature;
 
  SLIC_ERROR_IF(domain.type_ != Domain::Type::BoundaryElements,
                "Error: trying to evaluate a boundary integral over a non-boundary domain of integration");
 
  Integral integral(domain, argument_indices);
 
  if constexpr (dim == 1) {
    generate_bdr_kernels<mfem::Geometry::SEGMENT, Q>(signature, integral, qf);
  }
 
  if constexpr (dim == 2) {
    generate_bdr_kernels<mfem::Geometry::TRIANGLE, Q>(signature, integral, qf);
    generate_bdr_kernels<mfem::Geometry::SQUARE, Q>(signature, integral, qf);
  }
 
  return integral;
}
 
template <mfem::Geometry::Type geom, int Q, typename test, typename... trials, typename lambda_type>
void generate_interior_face_kernels(FunctionSignature<test(trials...)> s, Integral& integral, const lambda_type& qf)
{
  integral.geometric_factors_[geom] = GeometricFactors(integral.domain_, Q, geom);
  GeometricFactors& gf = integral.geometric_factors_[geom];
  if (gf.num_elements == 0) return;
 
  const double* positions = gf.X.Read();
  const double* jacobians = gf.J.Read();
  const uint32_t num_elements = uint32_t(gf.num_elements);
  const uint32_t qpts_per_element = num_quadrature_points(geom, Q);
 
  std::shared_ptr<zero> dummy_derivatives;
  integral.evaluation_[geom] = interior_face_integral::evaluation_kernel<NO_DIFFERENTIATION, Q, geom>(
      s, qf, positions, jacobians, dummy_derivatives, num_elements);
 
  constexpr std::size_t num_args = s.num_args;
  [[maybe_unused]] static constexpr int dim = dimension_of(geom);
  for_constexpr<num_args>([&](auto index) {
    // allocate memory for the derivatives of the q-function at each quadrature point
    //
    // Note: ptrs' lifetime is managed in an unusual way! It is captured by-value in the
    // action_of_gradient functor below to augment the reference count, and extend its lifetime to match
    // that of the boundaryIntegral that allocated it.
    using derivative_type = decltype(interior_face_integral::get_derivative_type<index, dim, trials...>(qf));
    auto ptr = accelerator::make_shared_array<ExecutionSpace::CPU, derivative_type>(num_elements * qpts_per_element);
 
    integral.evaluation_with_AD_[index][geom] =
        interior_face_integral::evaluation_kernel<index, Q, geom>(s, qf, positions, jacobians, ptr, num_elements);
 
    integral.jvp_[index][geom] =
        interior_face_integral::jacobian_vector_product_kernel<index, Q, geom>(s, ptr, num_elements);
    integral.element_gradient_[index][geom] =
        interior_face_integral::element_gradient_kernel<index, Q, geom>(s, ptr, num_elements);
  });
}
 
template <typename s, int Q, int dim, typename lambda_type>
Integral MakeInteriorFaceIntegral(const Domain& domain, const lambda_type& qf, std::vector<uint32_t> argument_indices)
{
  FunctionSignature<s> signature;
 
  SLIC_ERROR_IF(domain.type_ != Domain::Type::InteriorBoundaryElements,
                "Error: trying to evaluate a boundary integral over a non-boundary domain of integration");
 
  Integral integral(domain, argument_indices);
 
  if constexpr (dim == 1) {
    generate_interior_face_kernels<mfem::Geometry::SEGMENT, Q>(signature, integral, qf);
  }
 
  if constexpr (dim == 2) {
    generate_interior_face_kernels<mfem::Geometry::TRIANGLE, Q>(signature, integral, qf);
    generate_interior_face_kernels<mfem::Geometry::SQUARE, Q>(signature, integral, qf);
  }
 
  return integral;
}
 
}  // namespace smith