DirichletBC.h

// Copyright (C) 2007-2021 Michal Habera, Anders Logg, Garth N. Wells
// and Jørgen S.Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Constant.h"
#include "DofMap.h"
#include "Function.h"
#include "FunctionSpace.h"
#include <array>
#include <concepts>
#include <dolfinx/common/types.h>
#include <functional>
#include <memory>
#include <span>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
 
namespace dolfinx::fem
{
 
std::vector<std::int32_t>
locate_dofs_topological(const mesh::Topology& topology, const DofMap& dofmap,
                        int dim, std::span<const std::int32_t> entities,
                        bool remote = true);
 
std::array<std::vector<std::int32_t>, 2> locate_dofs_topological(
    const mesh::Topology& topology,
    std::array<std::reference_wrapper<const DofMap>, 2> dofmaps, int dim,
    std::span<const std::int32_t> entities, bool remote = true);
 
template <std::floating_point T, typename U>
std::vector<std::int32_t> locate_dofs_geometrical(const FunctionSpace<T>& V,
                                                  U marker_fn)
{
  // FIXME: Calling V.tabulate_dof_coordinates() is very expensive,
  // especially when we usually want the boundary dofs only. Add
  // interface that computes dofs coordinates only for specified cell.
 
  assert(V.element());
  if (V.element()->is_mixed())
  {
    throw std::runtime_error(
        "Cannot locate dofs geometrically for mixed space. Use subspaces.");
  }
 
  // Compute dof coordinates
  const std::vector<T> dof_coordinates = V.tabulate_dof_coordinates(true);
 
  using cmdspan3x_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
      const T,
      MDSPAN_IMPL_STANDARD_NAMESPACE::extents<
          std::size_t, 3, MDSPAN_IMPL_STANDARD_NAMESPACE::dynamic_extent>>;
 
  // Compute marker for each dof coordinate
  cmdspan3x_t x(dof_coordinates.data(), 3, dof_coordinates.size() / 3);
  const std::vector<std::int8_t> marked_dofs = marker_fn(x);
 
  std::vector<std::int32_t> dofs;
  dofs.reserve(std::count(marked_dofs.begin(), marked_dofs.end(), true));
  for (std::size_t i = 0; i < marked_dofs.size(); ++i)
  {
    if (marked_dofs[i])
      dofs.push_back(i);
  }
 
  return dofs;
}
 
template <std::floating_point T, typename U>
std::array<std::vector<std::int32_t>, 2> locate_dofs_geometrical(
    const std::array<std::reference_wrapper<const FunctionSpace<T>>, 2>& V,
    U marker_fn)
{
  // FIXME: Calling V.tabulate_dof_coordinates() is very expensive,
  // especially when we usually want the boundary dofs only. Add
  // interface that computes dofs coordinates only for specified cell.
 
  // Get function spaces
  const FunctionSpace<T>& V0 = V.at(0).get();
  const FunctionSpace<T>& V1 = V.at(1).get();
 
  // Get mesh
  auto mesh = V0.mesh();
  assert(mesh);
  assert(V1.mesh());
  if (mesh != V1.mesh())
    throw std::runtime_error("Meshes are not the same.");
  const int tdim = mesh->topology()->dim();
 
  assert(V0.element());
  assert(V1.element());
  if (*V0.element() != *V1.element())
    throw std::runtime_error("Function spaces must have the same element.");
 
  // Compute dof coordinates
  const std::vector<T> dof_coordinates = V1.tabulate_dof_coordinates(true);
 
  using cmdspan3x_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
      const T,
      MDSPAN_IMPL_STANDARD_NAMESPACE::extents<
          std::size_t, 3, MDSPAN_IMPL_STANDARD_NAMESPACE::dynamic_extent>>;
 
  // Evaluate marker for each dof coordinate
  cmdspan3x_t x(dof_coordinates.data(), 3, dof_coordinates.size() / 3);
  const std::vector<std::int8_t> marked_dofs = marker_fn(x);
 
  // Get dofmaps
  std::shared_ptr<const DofMap> dofmap0 = V0.dofmap();
  assert(dofmap0);
  const int bs0 = dofmap0->bs();
  std::shared_ptr<const DofMap> dofmap1 = V1.dofmap();
  assert(dofmap1);
  const int bs1 = dofmap1->bs();
 
  const int element_bs = dofmap0->element_dof_layout().block_size();
  assert(element_bs == dofmap1->element_dof_layout().block_size());
 
  // Iterate over cells
  auto topology = mesh->topology();
  assert(topology);
  std::vector<std::array<std::int32_t, 2>> bc_dofs;
  for (int c = 0; c < topology->connectivity(tdim, 0)->num_nodes(); ++c)
  {
    // Get cell dofmaps
    auto cell_dofs0 = dofmap0->cell_dofs(c);
    auto cell_dofs1 = dofmap1->cell_dofs(c);
 
    // Loop over cell dofs and add to bc_dofs if marked.
    for (std::size_t i = 0; i < cell_dofs1.size(); ++i)
    {
      if (marked_dofs[cell_dofs1[i]])
      {
        // Unroll over blocks
        for (int k = 0; k < element_bs; ++k)
        {
          const int local_pos = element_bs * i + k;
          const std::div_t pos0 = std::div(local_pos, bs0);
          const std::div_t pos1 = std::div(local_pos, bs1);
          const std::int32_t dof_index0
              = bs0 * cell_dofs0[pos0.quot] + pos0.rem;
          const std::int32_t dof_index1
              = bs1 * cell_dofs1[pos1.quot] + pos1.rem;
          bc_dofs.push_back({dof_index0, dof_index1});
        }
      }
    }
  }
 
  // Remove duplicates
  std::sort(bc_dofs.begin(), bc_dofs.end());
  bc_dofs.erase(std::unique(bc_dofs.begin(), bc_dofs.end()), bc_dofs.end());
 
  // Copy to separate array
  std::array dofs = {std::vector<std::int32_t>(bc_dofs.size()),
                     std::vector<std::int32_t>(bc_dofs.size())};
  std::transform(bc_dofs.cbegin(), bc_dofs.cend(), dofs[0].begin(),
                 [](auto dof) { return dof[0]; });
  std::transform(bc_dofs.cbegin(), bc_dofs.cend(), dofs[1].begin(),
                 [](auto dof) { return dof[1]; });
 
  return dofs;
}
 
template <dolfinx::scalar T,
          std::floating_point U = dolfinx::scalar_value_type_t<T>>
class DirichletBC
{
private:
  std::size_t num_owned(const DofMap& dofmap,
                        std::span<const std::int32_t> dofs)
  {
    int bs = dofmap.index_map_bs();
    std::int32_t map_size = dofmap.index_map->size_local();
    std::int32_t owned_size = bs * map_size;
    auto it = std::lower_bound(dofs.begin(), dofs.end(), owned_size);
    return std::distance(dofs.begin(), it);
  }
 
  static std::vector<std::int32_t>
  unroll_dofs(std::span<const std::int32_t> dofs, int bs)
  {
    std::vector<std::int32_t> dofs_unrolled(bs * dofs.size());
    for (std::size_t i = 0; i < dofs.size(); ++i)
      for (int k = 0; k < bs; ++k)
        dofs_unrolled[bs * i + k] = bs * dofs[i] + k;
    return dofs_unrolled;
  }
 
public:
  template <typename S, typename X,
            typename
            = std::enable_if_t<std::is_convertible_v<S, T>
                               or std::is_convertible_v<S, std::span<const T>>>>
    requires std::is_convertible_v<std::remove_cvref_t<X>,
                                   std::vector<std::int32_t>>
  DirichletBC(const S& g, X&& dofs, std::shared_ptr<const FunctionSpace<U>> V)
      : DirichletBC(std::make_shared<Constant<T>>(g), dofs, V)
  {
  }
 
  template <typename X>
    requires std::is_convertible_v<std::remove_cvref_t<X>,
                                   std::vector<std::int32_t>>
  DirichletBC(std::shared_ptr<const Constant<T>> g, X&& dofs,
              std::shared_ptr<const FunctionSpace<U>> V)
      : _function_space(V), _g(g), _dofs0(std::forward<X>(dofs)),
        _owned_indices0(num_owned(*V->dofmap(), _dofs0))
  {
    assert(g);
    assert(V);
    if (g->shape.size() != V->value_shape().size())
    {
      throw std::runtime_error(
          "Rank mis-match between Constant and function space in DirichletBC");
    }
 
    if (g->value.size() != _function_space->dofmap()->bs())
    {
      throw std::runtime_error(
          "Creating a DirichletBC using a Constant is not supported when the "
          "Constant size is not equal to the block size of the constrained "
          "(sub-)space. Use a fem::Function to create the fem::DirichletBC.");
    }
 
    if (!V->element()->interpolation_ident())
    {
      throw std::runtime_error(
          "Constant can be used only with point-evaluation elements");
    }
 
    // Unroll _dofs0 if dofmap block size > 1
    if (const int bs = V->dofmap()->bs(); bs > 1)
    {
      _owned_indices0 *= bs;
      _dofs0 = unroll_dofs(_dofs0, bs);
    }
  }
 
  template <typename X>
    requires std::is_convertible_v<std::remove_cvref_t<X>,
                                   std::vector<std::int32_t>>
  DirichletBC(std::shared_ptr<const Function<T, U>> g, X&& dofs)
      : _function_space(g->function_space()), _g(g),
        _dofs0(std::forward<X>(dofs)),
        _owned_indices0(num_owned(*_function_space->dofmap(), _dofs0))
  {
    assert(_function_space);
 
    // Unroll _dofs0 if dofmap block size > 1
    if (const int bs = _function_space->dofmap()->bs(); bs > 1)
    {
      _owned_indices0 *= bs;
      _dofs0 = unroll_dofs(_dofs0, bs);
    }
  }
 
  template <typename X>
  DirichletBC(std::shared_ptr<const Function<T, U>> g, X&& V_g_dofs,
              std::shared_ptr<const FunctionSpace<U>> V)
      : _function_space(V), _g(g),
        _dofs0(std::forward<typename X::value_type>(V_g_dofs[0])),
        _dofs1_g(std::forward<typename X::value_type>(V_g_dofs[1])),
        _owned_indices0(num_owned(*_function_space->dofmap(), _dofs0))
  {
  }
 
  DirichletBC(const DirichletBC& bc) = default;
 
  DirichletBC(DirichletBC&& bc) = default;
 
  ~DirichletBC() = default;
 
  DirichletBC& operator=(const DirichletBC& bc) = default;
 
  DirichletBC& operator=(DirichletBC&& bc) = default;
 
  std::shared_ptr<const FunctionSpace<U>> function_space() const
  {
    return _function_space;
  }
 
  std::variant<std::shared_ptr<const Function<T, U>>,
               std::shared_ptr<const Constant<T>>>
  value() const
  {
    return _g;
  }
 
  std::pair<std::span<const std::int32_t>, std::int32_t> dof_indices() const
  {
    return {_dofs0, _owned_indices0};
  }
 
  void set(std::span<T> x, T scale = 1) const
  {
    if (std::holds_alternative<std::shared_ptr<const Function<T, U>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Function<T, U>>>(_g);
      assert(g);
      std::span<const T> values = g->x()->array();
      auto dofs1_g = _dofs1_g.empty() ? std::span(_dofs0) : std::span(_dofs1_g);
      std::int32_t x_size = x.size();
      for (std::size_t i = 0; i < _dofs0.size(); ++i)
      {
        if (_dofs0[i] < x_size)
        {
          assert(dofs1_g[i] < (std::int32_t)values.size());
          x[_dofs0[i]] = scale * values[dofs1_g[i]];
        }
      }
    }
    else if (std::holds_alternative<std::shared_ptr<const Constant<T>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Constant<T>>>(_g);
      std::vector<T> value = g->value;
      int bs = _function_space->dofmap()->bs();
      std::int32_t x_size = x.size();
      std::for_each(_dofs0.cbegin(), _dofs0.cend(),
                    [x_size, bs, scale, &value, &x](auto dof)
                    {
                      if (dof < x_size)
                        x[dof] = scale * value[dof % bs];
                    });
    }
  }
 
  void set(std::span<T> x, std::span<const T> x0, T scale = 1) const
  {
    if (std::holds_alternative<std::shared_ptr<const Function<T, U>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Function<T, U>>>(_g);
      assert(g);
      std::span<const T> values = g->x()->array();
      assert(x.size() <= x0.size());
      auto dofs1_g = _dofs1_g.empty() ? std::span(_dofs0) : std::span(_dofs1_g);
      std::int32_t x_size = x.size();
      for (std::size_t i = 0; i < _dofs0.size(); ++i)
      {
        if (_dofs0[i] < x_size)
        {
          assert(dofs1_g[i] < (std::int32_t)values.size());
          x[_dofs0[i]] = scale * (values[dofs1_g[i]] - x0[_dofs0[i]]);
        }
      }
    }
    else if (std::holds_alternative<std::shared_ptr<const Constant<T>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Constant<T>>>(_g);
      const std::vector<T>& value = g->value;
      std::int32_t bs = _function_space->dofmap()->bs();
      std::for_each(_dofs0.begin(), _dofs0.end(),
                    [&x, &x0, &value, scale, bs](auto dof)
                    {
                      if (dof < (std::int32_t)x.size())
                        x[dof] = scale * (value[dof % bs] - x0[dof]);
                    });
    }
  }
 
  void dof_values(std::span<T> values) const
  {
    if (std::holds_alternative<std::shared_ptr<const Function<T, U>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Function<T, U>>>(_g);
      assert(g);
      std::span<const T> g_values = g->x()->array();
      auto dofs1_g = _dofs1_g.empty() ? std::span(_dofs0) : std::span(_dofs1_g);
      for (std::size_t i = 0; i < dofs1_g.size(); ++i)
        values[_dofs0[i]] = g_values[dofs1_g[i]];
    }
    else if (std::holds_alternative<std::shared_ptr<const Constant<T>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Constant<T>>>(_g);
      assert(g);
      const std::vector<T>& g_value = g->value;
      const std::int32_t bs = _function_space->dofmap()->bs();
      for (std::size_t i = 0; i < _dofs0.size(); ++i)
        values[_dofs0[i]] = g_value[_dofs0[i] % bs];
    }
  }
 
  void mark_dofs(std::span<std::int8_t> markers) const
  {
    for (std::size_t i = 0; i < _dofs0.size(); ++i)
    {
      assert(_dofs0[i] < (std::int32_t)markers.size());
      markers[_dofs0[i]] = true;
    }
  }
 
private:
  // The function space (possibly a sub function space)
  std::shared_ptr<const FunctionSpace<U>> _function_space;
 
  // The function
  std::variant<std::shared_ptr<const Function<T, U>>,
               std::shared_ptr<const Constant<T>>>
      _g;
 
  // Dof indices (_dofs0) in _function_space and (_dofs1_g) in the
  // space of _g. _dofs1_g may be empty if _dofs0 can be re-used
  std::vector<std::int32_t> _dofs0, _dofs1_g;
 
  // The first _owned_indices in _dofs are owned by this process
  std::int32_t _owned_indices0 = -1;
};
} // namespace dolfinx::fem