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 "Function.h"
#include "FunctionSpace.h"
#include <array>
#include <functional>
#include <memory>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
#include <xtensor/xtensor.hpp>
#include <xtl/xspan.hpp>
 
namespace dolfinx::fem
{
 
std::vector<std::int32_t>
locate_dofs_topological(const FunctionSpace& V, int dim,
                        const xtl::span<const std::int32_t>& entities,
                        bool remote = true);
 
std::array<std::vector<std::int32_t>, 2> locate_dofs_topological(
    const std::array<std::reference_wrapper<const FunctionSpace>, 2>& V,
    int dim, const xtl::span<const std::int32_t>& entities, bool remote = true);
 
std::vector<std::int32_t> locate_dofs_geometrical(
    const FunctionSpace& V,
    const std::function<xt::xtensor<bool, 1>(const xt::xtensor<double, 2>&)>&
        marker_fn);
 
std::array<std::vector<std::int32_t>, 2> locate_dofs_geometrical(
    const std::array<std::reference_wrapper<const FunctionSpace>, 2>& V,
    const std::function<xt::xtensor<bool, 1>(const xt::xtensor<double, 2>&)>&
        marker_fn);
 
template <typename T>
class DirichletBC
{
private:
  template <typename U>
  DirichletBC(const std::variant<std::shared_ptr<const Function<T>>,
                                 std::shared_ptr<const Constant<T>>>& g,
              U&& dofs, const std::shared_ptr<const FunctionSpace>& V, void*)
      : _function_space(V), _g(g), _dofs0(std::forward<U>(dofs))
  {
    assert(_function_space);
    if (auto c = std::get_if<std::shared_ptr<const Constant<T>>>(&_g))
    {
      assert(*c);
      if ((*c)->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 Function to create the DirichletBC.");
      }
    }
 
    // Compute number of owned dofs indices in the full space (will
    // contain 'gaps' for sub-spaces)
    const int map0_bs = _function_space->dofmap()->index_map_bs();
    const int map0_size = _function_space->dofmap()->index_map->size_local();
    const int owned_size0 = map0_bs * map0_size;
 
    // Find number of owned indices in _dofs0
    auto it0 = std::lower_bound(_dofs0.begin(), _dofs0.end(), owned_size0);
    _owned_indices0 = std::distance(_dofs0.begin(), it0);
 
    // Unroll _dofs0 for dofmap block size > 1
    if (const int bs = _function_space->dofmap()->bs(); bs > 1)
    {
      _owned_indices0 *= bs;
      const std::vector<std::int32_t> dof_tmp = _dofs0;
      _dofs0.resize(bs * dof_tmp.size());
      for (std::size_t i = 0; i < dof_tmp.size(); ++i)
      {
        for (int k = 0; k < bs; ++k)
          _dofs0[bs * i + k] = bs * dof_tmp[i] + k;
      }
    }
  }
 
public:
  template <typename S, typename U,
            typename = std::enable_if_t<
                std::is_convertible_v<
                    S, T> or std::is_convertible_v<S, xt::xarray<T>>>>
  DirichletBC(const S& g, U&& dofs,
              const std::shared_ptr<const FunctionSpace>& V)
      : DirichletBC(std::make_shared<Constant<T>>(g), dofs, V)
  {
  }
 
  template <typename U>
  DirichletBC(const std::shared_ptr<const Constant<T>>& g, U&& dofs,
              const std::shared_ptr<const FunctionSpace>& V)
      : DirichletBC(g, dofs, V, nullptr)
  {
    assert(g);
    assert(V);
    if (V->element()->value_shape().size() != g->shape.size())
    {
      throw std::runtime_error(
          "Rank mis-match between Constant and function space in DirichletBC");
    }
 
    if (!V->element()->interpolation_ident())
    {
      throw std::runtime_error(
          "Constant can be used only with point-evaluation elements");
    }
  }
 
  template <typename U>
  DirichletBC(const std::shared_ptr<const Function<T>>& g, U&& dofs)
      : DirichletBC(g, dofs, g->function_space(), nullptr)
  {
  }
 
  template <typename U>
  DirichletBC(const std::shared_ptr<const Function<T>>& g, U&& V_g_dofs,
              const std::shared_ptr<const FunctionSpace>& V)
      : _function_space(V), _g(g),
        _dofs0(std::forward<typename U::value_type>(V_g_dofs[0])),
        _dofs1_g(std::forward<typename U::value_type>(V_g_dofs[1]))
  {
    assert(_dofs0.size() == _dofs1_g.size());
    assert(_function_space);
    const int map0_bs = _function_space->dofmap()->index_map_bs();
    const int map0_size = _function_space->dofmap()->index_map->size_local();
    const int owned_size0 = map0_bs * map0_size;
    auto it0 = std::lower_bound(_dofs0.begin(), _dofs0.end(), owned_size0);
    _owned_indices0 = std::distance(_dofs0.begin(), it0);
  }
 
  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 fem::FunctionSpace> function_space() const
  {
    return _function_space;
  }
 
  std::variant<std::shared_ptr<const Function<T>>,
               std::shared_ptr<const Constant<T>>>
  value() const
  {
    return _g;
  }
 
  std::pair<xtl::span<const std::int32_t>, std::int32_t> dof_indices() const
  {
    return {_dofs0, _owned_indices0};
  }
 
  void set(xtl::span<T> x, double scale = 1.0) const
  {
    if (std::holds_alternative<std::shared_ptr<const Function<T>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Function<T>>>(_g);
      assert(g);
      xtl::span<const T> values = g->x()->array();
      auto dofs1_g = _dofs1_g.empty() ? xtl::span(_dofs0) : xtl::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(xtl::span<T> x, const xtl::span<const T>& x0,
           double scale = 1.0) const
  {
    if (std::holds_alternative<std::shared_ptr<const Function<T>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Function<T>>>(_g);
      assert(g);
      xtl::span<const T> values = g->x()->array();
      assert(x.size() <= x0.size());
      auto dofs1_g = _dofs1_g.empty() ? xtl::span(_dofs0) : xtl::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.cbegin(), _dofs0.cend(),
                    [&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(xtl::span<T> values) const
  {
    if (std::holds_alternative<std::shared_ptr<const Function<T>>>(_g))
    {
      auto g = std::get<std::shared_ptr<const Function<T>>>(_g);
      assert(g);
      xtl::span<const T> g_values = g->x()->array();
      auto dofs1_g = _dofs1_g.empty() ? xtl::span(_dofs0) : xtl::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(const xtl::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> _function_space;
 
  // The function
  std::variant<std::shared_ptr<const Function<T>>,
               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
  int _owned_indices0 = -1;
  int _owned_indices1 = -1;
};
} // namespace dolfinx::fem