DirichletBC.h

// Copyright (C) 2007-2020 Michal Habera, Anders Logg and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <dolfinx/fem/Function.h>
#include <dolfinx/fem/FunctionSpace.h>
#include <dolfinx/la/utils.h>
#include <functional>
#include <memory>
#include <vector>
#include <xtensor/xtensor.hpp>
#include <xtl/xspan.hpp>
 
namespace dolfinx::mesh
{
class Mesh;
}
 
namespace dolfinx::fem
{
 
std::array<std::vector<std::int32_t>, 2> locate_dofs_topological(
    const std::array<std::reference_wrapper<const fem::FunctionSpace>, 2>& V,
    const int dim, const xtl::span<const std::int32_t>& entities,
    bool remote = true);
 
std::vector<std::int32_t>
locate_dofs_topological(const fem::FunctionSpace& V, const int dim,
                        const xtl::span<const std::int32_t>& entities,
                        bool remote = true);
 
std::array<std::vector<std::int32_t>, 2> locate_dofs_geometrical(
    const std::array<std::reference_wrapper<const fem::FunctionSpace>, 2>& V,
    const std::function<xt::xtensor<bool, 1>(const xt::xtensor<double, 2>&)>&
        marker_fn);
 
std::vector<std::int32_t> locate_dofs_geometrical(
    const fem::FunctionSpace& V,
    const std::function<xt::xtensor<bool, 1>(const xt::xtensor<double, 2>&)>&
        marker_fn);
 
 
template <typename T>
class DirichletBC
{
 
public:
  template <typename U>
  DirichletBC(const std::shared_ptr<const fem::Function<T>>& g, U&& dofs)
      : _function_space(g->function_space()), _g(g),
        _dofs0(std::forward<U>(dofs))
  {
    const int owned_size0 = _function_space->dofmap()->index_map->size_local();
    auto it = std::lower_bound(_dofs0.begin(), _dofs0.end(), owned_size0);
    const int map0_bs = _function_space->dofmap()->index_map_bs();
    _owned_indices0 = map0_bs * std::distance(_dofs0.begin(), it);
 
    const int bs = _function_space->dofmap()->bs();
    if (bs > 1)
    {
      // Unroll for the block size
      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;
      }
    }
 
    // TODO: allows single dofs array (let one point to the other)
    _dofs1_g = _dofs0;
  }
 
  DirichletBC(const std::shared_ptr<const fem::Function<T>>& g,
              const std::array<std::vector<std::int32_t>, 2>& V_g_dofs,
              std::shared_ptr<const fem::FunctionSpace> V)
      : _function_space(V), _g(g), _dofs0(V_g_dofs[0]), _dofs1_g(V_g_dofs[1])
  {
    assert(_dofs0.size() == _dofs1_g.size());
    assert(_function_space);
    assert(_g);
 
    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::shared_ptr<const fem::Function<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
  {
    assert(_g);
    const std::vector<T>& g = _g->x()->array();
    for (std::size_t i = 0; i < _dofs0.size(); ++i)
    {
      if (_dofs0[i] < (std::int32_t)x.size())
      {
        assert(_dofs1_g[i] < (std::int32_t)g.size());
        x[_dofs0[i]] = scale * g[_dofs1_g[i]];
      }
    }
  }
 
  void set(xtl::span<T> x, const xtl::span<const T>& x0,
           double scale = 1.0) const
  {
    assert(_g);
    const std::vector<T>& g = _g->x()->array();
    assert(x.size() <= x0.size());
    for (std::size_t i = 0; i < _dofs0.size(); ++i)
    {
      if (_dofs0[i] < (std::int32_t)x.size())
      {
        assert(_dofs1_g[i] < (std::int32_t)g.size());
        x[_dofs0[i]] = scale * (g[_dofs1_g[i]] - x0[_dofs0[i]]);
      }
    }
  }
 
  void dof_values(xtl::span<T> values) const
  {
    assert(_g);
    const std::vector<T>& g = _g->x()->array();
    for (std::size_t i = 0; i < _dofs1_g.size(); ++i)
      values[_dofs0[i]] = g[_dofs1_g[i]];
  }
 
  void mark_dofs(std::vector<bool>& 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 fem::FunctionSpace> _function_space;
 
  // The function
  std::shared_ptr<const fem::Function<T>> _g;
 
  // Dof indices (_dofs0) in _function_space and (_dofs1_g) in the
  // space of _g
  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