Form.h

// Copyright (C) 2019-2020 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 <dolfinx/fem/FunctionSpace.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/MeshTags.h>
#include <functional>
#include <memory>
#include <string>
#include <vector>
 
namespace dolfinx::fem
{
 
template <typename T>
class Constant;
template <typename T>
class Function;
 
enum class IntegralType : std::int8_t
{
  cell = 0,
  exterior_facet = 1,
  interior_facet = 2,
  vertex = 3
};
 
 
template <typename T>
class Form
{
public:
  Form(const std::vector<std::shared_ptr<const fem::FunctionSpace>>&
           function_spaces,
       const std::map<
           IntegralType,
           std::pair<
               std::vector<std::pair<
                   int, std::function<void(
                            T*, const T*, const T*, const double*, const int*,
                            const std::uint8_t*, const std::uint32_t)>>>,
               const mesh::MeshTags<int>*>>& integrals,
       const std::vector<std::shared_ptr<const fem::Function<T>>>& coefficients,
       const std::vector<std::shared_ptr<const fem::Constant<T>>>& constants,
       bool needs_permutation_data,
       const std::shared_ptr<const mesh::Mesh>& mesh = nullptr)
      : _function_spaces(function_spaces), _coefficients(coefficients),
        _constants(constants), _mesh(mesh),
        _needs_permutation_data(needs_permutation_data)
  {
    // Extract _mesh from fem::FunctionSpace, and check they are the same
    if (!_mesh and !function_spaces.empty())
      _mesh = function_spaces[0]->mesh();
    for (const auto& V : function_spaces)
      if (_mesh != V->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& integral_type : integrals)
    {
      // Add key to map
      const IntegralType type = integral_type.first;
      auto it = _integrals.emplace(
          type, std::map<int, std::pair<kern, std::vector<std::int32_t>>>());
 
      // Loop over integrals kernels
      for (auto& integral : integral_type.second.first)
        it.first->second.insert({integral.first, {integral.second, {}}});
 
      // FIXME: do this neatly via a static function
      // Set domains for integral type
      if (integral_type.second.second)
      {
        assert(_mesh == integral_type.second.second->mesh());
        set_domains(type, *integral_type.second.second);
      }
    }
 
    // FIXME: do this neatly via a static function
    // Set markers for default integrals
    set_default_domains(*_mesh);
  }
 
  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> mesh() const { return _mesh; }
 
  const std::vector<std::shared_ptr<const fem::FunctionSpace>>&
  function_spaces() const
  {
    return _function_spaces;
  }
 
  const std::function<void(T*, const T*, const T*, const double*, const int*,
                           const std::uint8_t*, const std::uint32_t)>&
  kernel(IntegralType type, int i) const
  {
    auto it0 = _integrals.find(type);
    if (it0 == _integrals.end())
      throw std::runtime_error("No kernels for requested type.");
    auto it1 = it0->second.find(i);
    if (it1 == it0->second.end())
      throw std::runtime_error("No kernel for requested domain index.");
 
    return it1->second.first;
  }
 
  std::set<IntegralType> integral_types() const
  {
    std::set<IntegralType> set;
    for (auto& type : _integrals)
      set.insert(type.first);
    return set;
  }
 
  int num_integrals(IntegralType type) const
  {
    if (auto it = _integrals.find(type); it == _integrals.end())
      return 0;
    else
      return it->second.size();
  }
 
  std::vector<int> integral_ids(IntegralType type) const
  {
    std::vector<int> ids;
    if (auto it = _integrals.find(type); it != _integrals.end())
    {
      for (auto& kernel : it->second)
        ids.push_back(kernel.first);
    }
    return ids;
  }
 
  const std::vector<std::int32_t>& domains(IntegralType type, int i) const
  {
    auto it0 = _integrals.find(type);
    if (it0 == _integrals.end())
      throw std::runtime_error("No mesh entities for requested type.");
    auto it1 = it0->second.find(i);
    if (it1 == it0->second.end())
      throw std::runtime_error("No mesh entities for requested domain index.");
    return it1->second.second;
  }
 
  const std::vector<std::shared_ptr<const fem::Function<T>>>
  coefficients() const
  {
    return _coefficients;
  }
 
  bool needs_permutation_data() const { return _needs_permutation_data; }
 
  std::vector<int> coefficient_offsets() const
  {
    std::vector<int> n{0};
    for (const 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 fem::Constant<T>>>& constants() const
  {
    return _constants;
  }
 
  using scalar_type = T;
 
private:
  void set_domains(IntegralType type, const mesh::MeshTags<int>& marker)
  {
    auto it0 = _integrals.find(type);
    assert(it0 != _integrals.end());
 
    std::shared_ptr<const mesh::Mesh> mesh = marker.mesh();
    const mesh::Topology& topology = mesh->topology();
    const int tdim = topology.dim();
    int dim = tdim;
    if (type == IntegralType::exterior_facet
        or type == IntegralType::interior_facet)
    {
      dim = tdim - 1;
      mesh->topology_mutable().create_connectivity(dim, tdim);
    }
    else if (type == IntegralType::vertex)
      dim = 0;
 
    if (dim != marker.dim())
    {
      throw std::runtime_error("Invalid MeshTags dimension:"
                               + std::to_string(marker.dim()));
    }
 
    // Get all integrals for considered entity type
    std::map<int, std::pair<kern, std::vector<std::int32_t>>>& integrals
        = it0->second;
 
    // Get mesh tag data
    const std::vector<int>& values = marker.values();
    const std::vector<std::int32_t>& tagged_entities = marker.indices();
    assert(topology.index_map(dim));
    const auto entity_end
        = std::lower_bound(tagged_entities.begin(), tagged_entities.end(),
                           topology.index_map(dim)->size_local());
 
    if (dim == tdim - 1)
    {
      auto f_to_c = topology.connectivity(tdim - 1, tdim);
      assert(f_to_c);
      if (type == IntegralType::exterior_facet)
      {
        // Only need to consider shared facets when there are no ghost
        // cells
        assert(topology.index_map(tdim));
        std::set<std::int32_t> fwd_shared;
        if (topology.index_map(tdim)->num_ghosts() == 0)
        {
          fwd_shared.insert(
              topology.index_map(tdim - 1)->shared_indices().array().begin(),
              topology.index_map(tdim - 1)->shared_indices().array().end());
        }
 
        for (auto f = tagged_entities.begin(); f != entity_end; ++f)
        {
          // All "owned" facets connected to one cell, that are not
          // shared, should be external
          if (f_to_c->num_links(*f) == 1
              and fwd_shared.find(*f) == fwd_shared.end())
          {
            const std::size_t i = std::distance(tagged_entities.cbegin(), f);
            if (auto it = integrals.find(values[i]); it != integrals.end())
              it->second.second.push_back(*f);
          }
        }
      }
      else if (type == IntegralType::interior_facet)
      {
        for (auto f = tagged_entities.begin(); f != entity_end; ++f)
        {
          if (f_to_c->num_links(*f) == 2)
          {
            const std::size_t i = std::distance(tagged_entities.cbegin(), f);
            if (auto it = integrals.find(values[i]); it != integrals.end())
              it->second.second.push_back(*f);
          }
        }
      }
    }
    else
    {
      // For cell and vertex integrals use all markers (but not on ghost
      // entities)
      for (auto e = tagged_entities.begin(); e != entity_end; ++e)
      {
        const std::size_t i = std::distance(tagged_entities.cbegin(), e);
        if (auto it = integrals.find(values[i]); it != integrals.end())
          it->second.second.push_back(*e);
      }
    }
  }
 
  void set_default_domains(const mesh::Mesh& mesh)
  {
    const mesh::Topology& topology = mesh.topology();
    const int tdim = topology.dim();
 
    // Cells. If there is a default integral, define it on all owned cells
    if (auto kernels = _integrals.find(IntegralType::cell);
        kernels != _integrals.end())
    {
      if (auto it = kernels->second.find(-1); it != kernels->second.end())
      {
        std::vector<std::int32_t>& active_entities = it->second.second;
        const int num_cells = topology.index_map(tdim)->size_local();
        active_entities.resize(num_cells);
        std::iota(active_entities.begin(), active_entities.end(), 0);
      }
    }
 
    // Exterior facets. If there is a default integral, define it only
    // on owned surface facets.
    if (auto kernels = _integrals.find(IntegralType::exterior_facet);
        kernels != _integrals.end())
    {
      if (auto it = kernels->second.find(-1); it != kernels->second.end())
      {
        std::vector<std::int32_t>& active_entities = it->second.second;
        active_entities.clear();
 
        // Get number of facets owned by this process
        mesh.topology_mutable().create_connectivity(tdim - 1, tdim);
        auto f_to_c = topology.connectivity(tdim - 1, tdim);
        assert(topology.index_map(tdim - 1));
        std::set<std::int32_t> fwd_shared_facets;
 
        // Only need to consider shared facets when there are no ghost cells
        if (topology.index_map(tdim)->num_ghosts() == 0)
        {
          fwd_shared_facets.insert(
              topology.index_map(tdim - 1)->shared_indices().array().begin(),
              topology.index_map(tdim - 1)->shared_indices().array().end());
        }
 
        const int num_facets = topology.index_map(tdim - 1)->size_local();
        for (int f = 0; f < num_facets; ++f)
        {
          if (f_to_c->num_links(f) == 1
              and fwd_shared_facets.find(f) == fwd_shared_facets.end())
          {
            active_entities.push_back(f);
          }
        }
      }
    }
 
    // Interior facets. If there is a default integral, define it only on
    // owned interior facets.
    if (auto kernels = _integrals.find(IntegralType::interior_facet);
        kernels != _integrals.end())
    {
      if (auto it = kernels->second.find(-1); it != kernels->second.end())
      {
        std::vector<std::int32_t>& active_entities = it->second.second;
 
        // Get number of facets owned by this process
        mesh.topology_mutable().create_connectivity(tdim - 1, tdim);
        assert(topology.index_map(tdim - 1));
        const int num_facets = topology.index_map(tdim - 1)->size_local();
        auto f_to_c = topology.connectivity(tdim - 1, tdim);
        active_entities.clear();
        active_entities.reserve(num_facets);
        for (int f = 0; f < num_facets; ++f)
        {
          if (f_to_c->num_links(f) == 2)
            active_entities.push_back(f);
        }
      }
    }
  }
 
  // Function spaces (one for each argument)
  std::vector<std::shared_ptr<const fem::FunctionSpace>> _function_spaces;
 
  // Form coefficients
  std::vector<std::shared_ptr<const fem::Function<T>>> _coefficients;
 
  // Constants associated with the Form
  std::vector<std::shared_ptr<const fem::Constant<T>>> _constants;
 
  // The mesh
  std::shared_ptr<const mesh::Mesh> _mesh;
 
  using kern
      = std::function<void(T*, const T*, const T*, const double*, const int*,
                           const std::uint8_t*, const std::uint32_t)>;
  std::map<IntegralType,
           std::map<int, std::pair<kern, std::vector<std::int32_t>>>>
      _integrals;
 
  // True if permutation data needs to be passed into these integrals
  bool _needs_permutation_data;
};
} // namespace dolfinx::fem