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)
{
  return fem::impl::assemble_scalar(M);
}
 
// -- Vectors ----------------------------------------------------------------
 
template <typename T>
void assemble_vector(xtl::span<T> b, const Form<T>& L)
{
  fem::impl::assemble_vector(b, L);
}
 
// 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<std::vector<std::shared_ptr<const DirichletBC<T>>>>& bcs1,
    const std::vector<xtl::span<const T>>& x0, double scale)
{
  fem::impl::apply_lifting(b, a, bcs1, x0, scale);
}
 
// -- Matrices ---------------------------------------------------------------
 
// Experimental
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)
{
  // 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, 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)
 
{
  impl::assemble_matrix(mat_add, a, 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