Form.h

// Copyright (C) 2019-2023 Garth N. Wells and Chris Richardson
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "FunctionSpace.h"
#include <algorithm>
#include <array>
#include <concepts>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/Mesh.h>
#include <functional>
#include <memory>
#include <span>
#include <string>
#include <tuple>
#include <vector>
 
namespace dolfinx::fem
{
 
template <dolfinx::scalar T>
class Constant;
template <dolfinx::scalar T, std::floating_point U>
class Function;
 
enum class IntegralType : std::int8_t
{
  cell = 0,           
  exterior_facet = 1, 
  interior_facet = 2, 
  vertex = 3          
};
 
template <dolfinx::scalar T,
          std::floating_point U = dolfinx::scalar_value_type_t<T>>
class Form
{
public:
  using scalar_type = T;
 
  Form(const std::vector<std::shared_ptr<const FunctionSpace<U>>>& V,
       const std::map<IntegralType,
                      std::vector<std::tuple<
                          int,
                          std::function<void(T*, const T*, const T*,
                                             const scalar_value_type_t<T>*,
                                             const int*, const std::uint8_t*)>,
                          std::vector<std::int32_t>>>>& integrals,
       const std::vector<std::shared_ptr<const Function<T, U>>>& coefficients,
       const std::vector<std::shared_ptr<const Constant<T>>>& constants,
       bool needs_facet_permutations,
       std::shared_ptr<const mesh::Mesh<U>> mesh = nullptr)
      : _function_spaces(V), _coefficients(coefficients), _constants(constants),
        _mesh(mesh), _needs_facet_permutations(needs_facet_permutations)
  {
    // Extract _mesh from FunctionSpace, and check they are the same
    if (!_mesh and !V.empty())
      _mesh = V[0]->mesh();
    for (auto& space : V)
    {
      if (_mesh != space->mesh())
        throw std::runtime_error("Incompatible mesh");
    }
    if (!_mesh)
      throw std::runtime_error("No mesh could be associated with the Form.");
 
    // Store kernels, looping over integrals by domain type (dimension)
    for (auto& [type, kernels] : integrals)
    {
      auto& integrals = _integrals[static_cast<std::size_t>(type)];
      for (auto& [id, kern, e] : kernels)
        integrals.insert({id, {kern, std::vector(e.begin(), e.end())}});
    }
  }
 
  Form(const Form& form) = delete;
 
  Form(Form&& form) = default;
 
  virtual ~Form() = default;
 
  int rank() const { return _function_spaces.size(); }
 
  std::shared_ptr<const mesh::Mesh<U>> mesh() const { return _mesh; }
 
  const std::vector<std::shared_ptr<const FunctionSpace<U>>>&
  function_spaces() const
  {
    return _function_spaces;
  }
 
  std::function<void(T*, const T*, const T*, const scalar_value_type_t<T>*,
                     const int*, const std::uint8_t*)>
  kernel(IntegralType type, int i) const
  {
    auto integrals = _integrals[static_cast<std::size_t>(type)];
    if (auto it = integrals.find(i); it != integrals.end())
      return it->second.first;
    else
      throw std::runtime_error("No kernel for requested domain index.");
  }
 
  std::set<IntegralType> integral_types() const
  {
    std::set<IntegralType> set;
    for (std::size_t i = 0; i < _integrals.size(); ++i)
    {
      if (!_integrals[i].empty())
        set.insert(static_cast<IntegralType>(i));
    }
 
    return set;
  }
 
  int num_integrals(IntegralType type) const
  {
    return _integrals[static_cast<std::size_t>(type)].size();
  }
 
  std::vector<int> integral_ids(IntegralType type) const
  {
    std::vector<int> ids;
    auto& integrals = _integrals[static_cast<std::size_t>(type)];
    std::transform(integrals.begin(), integrals.end(), std::back_inserter(ids),
                   [](auto& integral) { return integral.first; });
    return ids;
  }
 
  std::span<const std::int32_t> domain(IntegralType type, int i) const
  {
    auto& integral = _integrals[static_cast<std::size_t>(type)];
    if (auto it = integral.find(i); it != integral.end())
      return it->second.second;
    else
      throw std::runtime_error("No mesh entities for requested domain index.");
  }
 
  const std::vector<std::shared_ptr<const Function<T, U>>>& coefficients() const
  {
    return _coefficients;
  }
 
  bool needs_facet_permutations() const { return _needs_facet_permutations; }
 
  std::vector<int> coefficient_offsets() const
  {
    std::vector<int> n = {0};
    for (auto& c : _coefficients)
    {
      if (!c)
        throw std::runtime_error("Not all form coefficients have been set.");
      n.push_back(n.back() + c->function_space()->element()->space_dimension());
    }
    return n;
  }
 
  const std::vector<std::shared_ptr<const Constant<T>>>& constants() const
  {
    return _constants;
  }
 
private:
  using kern = std::function<void(T*, const T*, const T*,
                                  const scalar_value_type_t<T>*, const int*,
                                  const std::uint8_t*)>;
 
  // Function spaces (one for each argument)
  std::vector<std::shared_ptr<const FunctionSpace<U>>> _function_spaces;
 
  // Form coefficients
  std::vector<std::shared_ptr<const Function<T, U>>> _coefficients;
 
  // Constants associated with the Form
  std::vector<std::shared_ptr<const Constant<T>>> _constants;
 
  // The mesh
  std::shared_ptr<const mesh::Mesh<U>> _mesh;
 
  // Integrals. Array index is
  // static_cast<std::size_t(IntegralType::foo)
  std::array<std::map<int, std::pair<kern, std::vector<std::int32_t>>>, 4>
      _integrals;
 
  // True if permutation data needs to be passed into these integrals
  bool _needs_facet_permutations;
};
} // namespace dolfinx::fem