assembler.h

// Copyright (C) 2018-2020 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "assemble_matrix_impl.h"
#include "assemble_scalar_impl.h"
#include "assemble_vector_impl.h"
#include <memory>
#include <vector>
#include <xtl/xspan.hpp>
 
namespace dolfinx::fem
{
template <typename T>
class DirichletBC;
template <typename T>
class Form;
class FunctionSpace;
 
// -- Scalar ----------------------------------------------------------------
 
template <typename T>
T assemble_scalar(const Form<T>& M, const xtl::span<const T>& constants,
                  const array2d<T>& coeffs)
{
  return impl::assemble_scalar(M, constants, coeffs);
}
 
template <typename T>
T assemble_scalar(const Form<T>& M)
{
  const std::vector<T> constants = pack_constants(M);
  const array2d<T> coeffs = pack_coefficients(M);
  return assemble_scalar(M, tcb::make_span(constants), coeffs);
}
 
// -- Vectors ----------------------------------------------------------------
 
template <typename T>
void assemble_vector(xtl::span<T> b, const Form<T>& L,
                     const xtl::span<const T>& constants,
                     const array2d<T>& coeffs)
{
  impl::assemble_vector(b, L, constants, coeffs);
}
 
template <typename T>
void assemble_vector(xtl::span<T> b, const Form<T>& L)
{
  const std::vector<T> constants = pack_constants(L);
  const array2d<T> coeffs = pack_coefficients(L);
  assemble_vector(b, L, tcb::make_span(constants), coeffs);
}
 
// FIXME: clarify how x0 is used
// FIXME: if bcs entries are set
 
// FIXME: need to pass an array of Vec for x0?
// FIXME: clarify zeroing of vector
 
template <typename T>
void apply_lifting(
    xtl::span<T> b, const std::vector<std::shared_ptr<const Form<T>>>& a,
    const std::vector<xtl::span<const T>>& constants,
    const std::vector<const array2d<T>*>& coeffs,
    const std::vector<std::vector<std::shared_ptr<const DirichletBC<T>>>>& bcs1,
    const std::vector<xtl::span<const T>>& x0, double scale)
{
  impl::apply_lifting(b, a, constants, coeffs, bcs1, x0, scale);
}
 
template <typename T>
void apply_lifting(
    xtl::span<T> b, const std::vector<std::shared_ptr<const Form<T>>>& a,
    const std::vector<std::vector<std::shared_ptr<const DirichletBC<T>>>>& bcs1,
    const std::vector<xtl::span<const T>>& x0, double scale)
{
  std::vector<std::unique_ptr<array2d<T>>> coeffs;
  std::vector<std::vector<T>> constants;
  for (auto _a : a)
  {
    if (_a)
    {
      coeffs.push_back(std::make_unique<array2d<T>>(pack_coefficients(*_a)));
      constants.push_back(pack_constants(*_a));
    }
    else
    {
      coeffs.push_back(nullptr);
      constants.push_back({});
    }
  }
 
  std::vector<const array2d<T>*> _coeffs;
  std::for_each(coeffs.begin(), coeffs.end(),
                [&_coeffs](const auto& c) { _coeffs.push_back(c.get()); });
  std::vector<xtl::span<const T>> _constants(constants.begin(),
                                             constants.end());
  apply_lifting(b, a, _constants, _coeffs, bcs1, x0, scale);
}
 
// -- Matrices ---------------------------------------------------------------
 
template <typename T>
void assemble_matrix(
    const std::function<int(std::int32_t, const std::int32_t*, std::int32_t,
                            const std::int32_t*, const T*)>& mat_add,
    const Form<T>& a, const xtl::span<const T>& constants,
    const array2d<T>& coeffs,
    const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs)
{
  // Index maps for dof ranges
  auto map0 = a.function_spaces().at(0)->dofmap()->index_map;
  auto map1 = a.function_spaces().at(1)->dofmap()->index_map;
  auto bs0 = a.function_spaces().at(0)->dofmap()->index_map_bs();
  auto bs1 = a.function_spaces().at(1)->dofmap()->index_map_bs();
 
  // Build dof markers
  std::vector<bool> dof_marker0, dof_marker1;
  assert(map0);
  std::int32_t dim0 = bs0 * (map0->size_local() + map0->num_ghosts());
  assert(map1);
  std::int32_t dim1 = bs1 * (map1->size_local() + map1->num_ghosts());
  for (std::size_t k = 0; k < bcs.size(); ++k)
  {
    assert(bcs[k]);
    assert(bcs[k]->function_space());
    if (a.function_spaces().at(0)->contains(*bcs[k]->function_space()))
    {
      dof_marker0.resize(dim0, false);
      bcs[k]->mark_dofs(dof_marker0);
    }
 
    if (a.function_spaces().at(1)->contains(*bcs[k]->function_space()))
    {
      dof_marker1.resize(dim1, false);
      bcs[k]->mark_dofs(dof_marker1);
    }
  }
 
  // Assemble
  impl::assemble_matrix(mat_add, a, constants, coeffs, dof_marker0,
                        dof_marker1);
}
 
template <typename T>
void assemble_matrix(
    const std::function<int(std::int32_t, const std::int32_t*, std::int32_t,
                            const std::int32_t*, const T*)>& mat_add,
    const Form<T>& a,
    const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs)
{
  // Prepare constants and coefficients
  const std::vector<T> constants = pack_constants(a);
  const array2d<T> coeffs = pack_coefficients(a);
 
  // Assemble
  assemble_matrix(mat_add, a, tcb::make_span(constants), coeffs, bcs);
}
 
template <typename T>
void assemble_matrix(
    const std::function<int(std::int32_t, const std::int32_t*, std::int32_t,
                            const std::int32_t*, const T*)>& mat_add,
    const Form<T>& a, const xtl::span<const T>& constants,
    const array2d<T>& coeffs, const std::vector<bool>& dof_marker0,
    const std::vector<bool>& dof_marker1)
 
{
  impl::assemble_matrix(mat_add, a, constants, coeffs, dof_marker0,
                        dof_marker1);
}
 
template <typename T>
void assemble_matrix(
    const std::function<int(std::int32_t, const std::int32_t*, std::int32_t,
                            const std::int32_t*, const T*)>& mat_add,
    const Form<T>& a, const std::vector<bool>& dof_marker0,
    const std::vector<bool>& dof_marker1)
 
{
  // Prepare constants and coefficients
  const std::vector<T> constants = pack_constants(a);
  const array2d<T> coeffs = pack_coefficients(a);
 
  // Assemble
  assemble_matrix(mat_add, a, tcb::make_span(constants), coeffs, dof_marker0,
                  dof_marker1);
}
 
template <typename T>
void set_diagonal(
    const std::function<int(std::int32_t, const std::int32_t*, std::int32_t,
                            const std::int32_t*, const T*)>& set_fn,
    const xtl::span<const std::int32_t>& rows, T diagonal = 1.0)
{
  for (std::size_t i = 0; i < rows.size(); ++i)
  {
    const std::int32_t row = rows[i];
    set_fn(1, &row, 1, &row, &diagonal);
  }
}
 
template <typename T>
void set_diagonal(
    const std::function<int(std::int32_t, const std::int32_t*, std::int32_t,
                            const std::int32_t*, const T*)>& set_fn,
    const fem::FunctionSpace& V,
    const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
    T diagonal = 1.0)
{
  for (const auto& bc : bcs)
  {
    assert(bc);
    if (V.contains(*bc->function_space()))
    {
      const auto [dofs, range] = bc->dof_indices();
      set_diagonal<T>(set_fn, dofs.first(range), diagonal);
    }
  }
}
 
// -- Setting bcs ------------------------------------------------------------
 
// FIXME: Move these function elsewhere?
 
// FIXME: clarify x0
// FIXME: clarify what happens with ghosts
 
template <typename T>
void set_bc(xtl::span<T> b,
            const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
            const xtl::span<const T>& x0, double scale = 1.0)
{
  if (b.size() > x0.size())
    throw std::runtime_error("Size mismatch between b and x0 vectors.");
  for (const auto& bc : bcs)
  {
    assert(bc);
    bc->set(b, x0, scale);
  }
}
 
template <typename T>
void set_bc(xtl::span<T> b,
            const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
            double scale = 1.0)
{
  for (const auto& bc : bcs)
  {
    assert(bc);
    bc->set(b, scale);
  }
}
 
// FIXME: Handle null block
// FIXME: Pass function spaces rather than forms
 
template <typename T>
std::vector<std::vector<std::shared_ptr<const fem::DirichletBC<T>>>>
bcs_rows(const std::vector<const Form<T>*>& L,
         const std::vector<std::shared_ptr<const fem::DirichletBC<T>>>& bcs)
{
  // Pack DirichletBC pointers for rows
  std::vector<std::vector<std::shared_ptr<const fem::DirichletBC<T>>>> bcs0(
      L.size());
  for (std::size_t i = 0; i < L.size(); ++i)
    for (const std::shared_ptr<const DirichletBC<T>>& bc : bcs)
      if (L[i]->function_spaces()[0]->contains(*bc->function_space()))
        bcs0[i].push_back(bc);
  return bcs0;
}
 
// FIXME: Handle null block
// FIXME: Pass function spaces rather than forms
 
template <typename T>
std::vector<
    std::vector<std::vector<std::shared_ptr<const fem::DirichletBC<T>>>>>
bcs_cols(const std::vector<std::vector<std::shared_ptr<const Form<T>>>>& a,
         const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs)
{
  // Pack DirichletBC pointers for columns
  std::vector<
      std::vector<std::vector<std::shared_ptr<const fem::DirichletBC<T>>>>>
      bcs1(a.size());
  for (std::size_t i = 0; i < a.size(); ++i)
  {
    for (std::size_t j = 0; j < a[i].size(); ++j)
    {
      bcs1[i].resize(a[j].size());
      for (const std::shared_ptr<const DirichletBC<T>>& bc : bcs)
      {
        // FIXME: handle case where a[i][j] is null
        if (a[i][j])
        {
          if (a[i][j]->function_spaces()[1]->contains(*bc->function_space()))
            bcs1[i][j].push_back(bc);
        }
      }
    }
  }
 
  return bcs1;
}
 
} // namespace dolfinx::fem