Loading...
Searching...
No Matches
assemble_matrix_impl.h
// Copyright (C) 2018-2019 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once
#include "DofMap.h"
#include "Form.h"
#include "FunctionSpace.h"
#include "traits.h"
#include "utils.h"
#include <algorithm>
#include <dolfinx/la/utils.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <functional>
#include <iterator>
#include <span>
#include <tuple>
#include <vector>
namespace dolfinx::fem::impl
{
using mdspan2_t = md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>;
template <dolfinx::scalar T>
void assemble_cells_matrix(
la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
md::mdspan<const scalar_value_t<T>,
md::extents<std::size_t, md::dynamic_extent, 3>>
x,
std::span<const std::int32_t> cells,
std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap0,
fem::DofTransformKernel<T> auto P0,
std::tuple<mdspan2_t, int, std::span<const std::int32_t>> dofmap1,
fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
std::span<const std::uint32_t> cell_info1)
{
if (cells.empty())
return;
const auto [dmap0, bs0, cells0] = dofmap0;
const auto [dmap1, bs1, cells1] = dofmap1;
// Iterate over active cells
const int num_dofs0 = dmap0.extent(1);
const int num_dofs1 = dmap1.extent(1);
const int ndim0 = bs0 * num_dofs0;
const int ndim1 = bs1 * num_dofs1;
std::vector<T> Ae(ndim0 * ndim1);
std::vector<scalar_value_t<T>> cdofs(3 * x_dofmap.extent(1));
// Iterate over active cells
assert(cells0.size() == cells.size());
assert(cells1.size() == cells.size());
for (std::size_t c = 0; c < cells.size(); ++c)
{
// Cell index in integration domain mesh (c), test function mesh
// (c0) and trial function mesh (c1)
std::int32_t cell = cells[c];
std::int32_t cell0 = cells0[c];
std::int32_t cell1 = cells1[c];
// Get cell coordinates/geometry
auto x_dofs = md::submdspan(x_dofmap, cell, md::full_extent);
for (std::size_t i = 0; i < x_dofs.size(); ++i)
std::copy_n(&x(x_dofs[i], 0), 3, std::next(cdofs.begin(), 3 * i));
// Tabulate tensor
std::ranges::fill(Ae, 0);
kernel(Ae.data(), &coeffs(c, 0), constants.data(), cdofs.data(), nullptr,
nullptr, nullptr);
// Compute A = P_0 \tilde{A} P_1^T (dof transformation)
P0(Ae, cell_info0, cell0, ndim1); // B = P0 \tilde{A}
P1T(Ae, cell_info1, cell1, ndim0); // A = B P1_T
// Zero rows/columns for essential bcs
std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
if (!bc0.empty())
{
for (int i = 0; i < num_dofs0; ++i)
{
for (int k = 0; k < bs0; ++k)
{
if (bc0[bs0 * dofs0[i] + k])
{
// Zero row bs0 * i + k
const int row = bs0 * i + k;
std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
}
}
}
}
if (!bc1.empty())
{
for (int j = 0; j < num_dofs1; ++j)
{
for (int k = 0; k < bs1; ++k)
{
if (bc1[bs1 * dofs1[j] + k])
{
// Zero column bs1 * j + k
const int col = bs1 * j + k;
for (int row = 0; row < ndim0; ++row)
Ae[row * ndim1 + col] = 0;
}
}
}
}
mat_set(dofs0, dofs1, Ae);
}
}
template <dolfinx::scalar T>
void assemble_entities(
la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
md::mdspan<const scalar_value_t<T>,
md::extents<std::size_t, md::dynamic_extent, 3>>
x,
md::mdspan<const std::int32_t,
std::extents<std::size_t, md::dynamic_extent, 2>>
entities,
std::tuple<mdspan2_t, int,
md::mdspan<const std::int32_t,
std::extents<std::size_t, md::dynamic_extent, 2>>>
dofmap0,
fem::DofTransformKernel<T> auto P0,
std::tuple<mdspan2_t, int,
md::mdspan<const std::int32_t,
std::extents<std::size_t, md::dynamic_extent, 2>>>
dofmap1,
fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
std::span<const std::uint32_t> cell_info1,
md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
{
if (entities.empty())
return;
const auto [dmap0, bs0, entities0] = dofmap0;
const auto [dmap1, bs1, entities1] = dofmap1;
// Data structures used in assembly
std::vector<scalar_value_t<T>> cdofs(3 * x_dofmap.extent(1));
const int num_dofs0 = dmap0.extent(1);
const int num_dofs1 = dmap1.extent(1);
const int ndim0 = bs0 * num_dofs0;
const int ndim1 = bs1 * num_dofs1;
std::vector<T> Ae(ndim0 * ndim1);
assert(entities0.size() == entities.size());
assert(entities1.size() == entities.size());
for (std::size_t f = 0; f < entities.extent(0); ++f)
{
// Cell in the integration domain, local entity index relative to the
// integration domain cell, and cells in the test and trial function
// meshes
std::int32_t cell = entities(f, 0);
std::int32_t local_entity = entities(f, 1);
std::int32_t cell0 = entities0(f, 0);
std::int32_t cell1 = entities1(f, 0);
// Get cell coordinates/geometry
auto x_dofs = md::submdspan(x_dofmap, cell, md::full_extent);
for (std::size_t i = 0; i < x_dofs.size(); ++i)
std::copy_n(&x(x_dofs[i], 0), 3, std::next(cdofs.begin(), 3 * i));
// Permutations
std::uint8_t perm = perms.empty() ? 0 : perms(cell, local_entity);
// Tabulate tensor
std::ranges::fill(Ae, 0);
kernel(Ae.data(), &coeffs(f, 0), constants.data(), cdofs.data(),
&local_entity, &perm, nullptr);
P0(Ae, cell_info0, cell0, ndim1);
P1T(Ae, cell_info1, cell1, ndim0);
// Zero rows/columns for essential bcs
std::span dofs0(dmap0.data_handle() + cell0 * num_dofs0, num_dofs0);
std::span dofs1(dmap1.data_handle() + cell1 * num_dofs1, num_dofs1);
if (!bc0.empty())
{
for (int i = 0; i < num_dofs0; ++i)
{
for (int k = 0; k < bs0; ++k)
{
if (bc0[bs0 * dofs0[i] + k])
{
// Zero row bs0 * i + k
const int row = bs0 * i + k;
std::fill_n(std::next(Ae.begin(), ndim1 * row), ndim1, 0);
}
}
}
}
if (!bc1.empty())
{
for (int j = 0; j < num_dofs1; ++j)
{
for (int k = 0; k < bs1; ++k)
{
if (bc1[bs1 * dofs1[j] + k])
{
// Zero column bs1 * j + k
const int col = bs1 * j + k;
for (int row = 0; row < ndim0; ++row)
Ae[row * ndim1 + col] = 0;
}
}
}
}
mat_set(dofs0, dofs1, Ae);
}
}
template <dolfinx::scalar T>
void assemble_interior_facets(
la::MatSet<T> auto mat_set, mdspan2_t x_dofmap,
md::mdspan<const scalar_value_t<T>,
md::extents<std::size_t, md::dynamic_extent, 3>>
x,
md::mdspan<const std::int32_t,
std::extents<std::size_t, md::dynamic_extent, 2, 2>>
facets,
std::tuple<const DofMap&, int,
md::mdspan<const std::int32_t,
std::extents<std::size_t, md::dynamic_extent, 2, 2>>>
dofmap0,
fem::DofTransformKernel<T> auto P0,
std::tuple<const DofMap&, int,
md::mdspan<const std::int32_t,
std::extents<std::size_t, md::dynamic_extent, 2, 2>>>
dofmap1,
fem::DofTransformKernel<T> auto P1T, std::span<const std::int8_t> bc0,
std::span<const std::int8_t> bc1, FEkernel<T> auto kernel,
md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
md::dynamic_extent>>
coeffs,
std::span<const T> constants, std::span<const std::uint32_t> cell_info0,
std::span<const std::uint32_t> cell_info1,
md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
{
if (facets.empty())
return;
const auto [dmap0, bs0, facets0] = dofmap0;
const auto [dmap1, bs1, facets1] = dofmap1;
// Data structures used in assembly
using X = scalar_value_t<T>;
std::vector<X> cdofs(2 * x_dofmap.extent(1) * 3);
std::span<X> cdofs0(cdofs.data(), x_dofmap.extent(1) * 3);
std::span<X> cdofs1(cdofs.data() + x_dofmap.extent(1) * 3,
x_dofmap.extent(1) * 3);
const std::size_t dmap0_size = dmap0.map().extent(1);
const std::size_t dmap1_size = dmap1.map().extent(1);
const int num_rows = bs0 * 2 * dmap0_size;
const int num_cols = bs1 * 2 * dmap1_size;
// Temporaries for joint dofmaps
std::vector<T> Ae(num_rows * num_cols), be(num_rows);
std::vector<std::int32_t> dmapjoint0(2 * dmap0_size);
std::vector<std::int32_t> dmapjoint1(2 * dmap1_size);
assert(facets0.size() == facets.size());
assert(facets1.size() == facets.size());
for (std::size_t f = 0; f < facets.extent(0); ++f)
{
// Cells in integration domain, test function domain and trial
// function domain
std::array cells{facets(f, 0, 0), facets(f, 1, 0)};
std::array cells0{facets0(f, 0, 0), facets0(f, 1, 0)};
std::array cells1{facets1(f, 0, 0), facets1(f, 1, 0)};
// Local facets indices
std::array local_facet{facets(f, 0, 1), facets(f, 1, 1)};
// Get cell geometry
auto x_dofs0 = md::submdspan(x_dofmap, cells[0], md::full_extent);
for (std::size_t i = 0; i < x_dofs0.size(); ++i)
std::copy_n(&x(x_dofs0[i], 0), 3, std::next(cdofs0.begin(), 3 * i));
auto x_dofs1 = md::submdspan(x_dofmap, cells[1], md::full_extent);
for (std::size_t i = 0; i < x_dofs1.size(); ++i)
std::copy_n(&x(x_dofs1[i], 0), 3, std::next(cdofs1.begin(), 3 * i));
// Get dof maps for cells and pack
// When integrating over interfaces between two domains, the test function
// might only be defined on one side, so we check which cells exist in the
// test function domain
std::span<const std::int32_t> dmap0_cell0
= cells0[0] >= 0 ? dmap0.cell_dofs(cells0[0])
: std::span<const std::int32_t>();
std::span<const std::int32_t> dmap0_cell1
= cells0[1] >= 0 ? dmap0.cell_dofs(cells0[1])
: std::span<const std::int32_t>();
std::ranges::copy(dmap0_cell0, dmapjoint0.begin());
std::ranges::copy(dmap0_cell1, std::next(dmapjoint0.begin(), dmap0_size));
// Check which cells exist in the trial function domain
std::span<const std::int32_t> dmap1_cell0
= cells1[0] >= 0 ? dmap1.cell_dofs(cells1[0])
: std::span<const std::int32_t>();
std::span<const std::int32_t> dmap1_cell1
= cells1[1] >= 0 ? dmap1.cell_dofs(cells1[1])
: std::span<const std::int32_t>();
std::ranges::copy(dmap1_cell0, dmapjoint1.begin());
std::ranges::copy(dmap1_cell1, std::next(dmapjoint1.begin(), dmap1_size));
// Tabulate tensor
std::ranges::fill(Ae, 0);
std::array perm = perms.empty()
? std::array<std::uint8_t, 2>{0, 0}
: std::array{perms(cells[0], local_facet[0]),
perms(cells[1], local_facet[1])};
kernel(Ae.data(), &coeffs(f, 0, 0), constants.data(), cdofs.data(),
local_facet.data(), perm.data(), nullptr);
// Local element layout is a 2x2 block matrix with structure
//
// cell0cell0 | cell0cell1
// cell1cell0 | cell1cell1
//
// where each block is element tensor of size (dmap0, dmap1).
// Only apply transformation when cells exist
if (cells0[0] >= 0)
P0(Ae, cell_info0, cells0[0], num_cols);
if (cells0[1] >= 0)
{
std::span sub_Ae0(Ae.data() + bs0 * dmap0_size * num_cols,
bs0 * dmap0_size * num_cols);
P0(sub_Ae0, cell_info0, cells0[1], num_cols);
}
if (cells1[0] >= 0)
P1T(Ae, cell_info1, cells1[0], num_rows);
if (cells1[1] >= 0)
{
for (int row = 0; row < num_rows; ++row)
{
// DOFs for dmap1 and cell1 are not stored contiguously in the
// block matrix, so each row needs a separate span access
std::span sub_Ae1(Ae.data() + row * num_cols + bs1 * dmap1_size,
bs1 * dmap1_size);
P1T(sub_Ae1, cell_info1, cells1[1], 1);
}
}
// Zero rows/columns for essential bcs
if (!bc0.empty())
{
for (std::size_t i = 0; i < dmapjoint0.size(); ++i)
{
for (int k = 0; k < bs0; ++k)
{
if (bc0[bs0 * dmapjoint0[i] + k])
{
// Zero row bs0 * i + k
std::fill_n(std::next(Ae.begin(), num_cols * (bs0 * i + k)),
num_cols, 0);
}
}
}
}
if (!bc1.empty())
{
for (std::size_t j = 0; j < dmapjoint1.size(); ++j)
{
for (int k = 0; k < bs1; ++k)
{
if (bc1[bs1 * dmapjoint1[j] + k])
{
// Zero column bs1 * j + k
for (int m = 0; m < num_rows; ++m)
Ae[m * num_cols + bs1 * j + k] = 0;
}
}
}
}
mat_set(dmapjoint0, dmapjoint1, Ae);
}
}
template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix(
la::MatSet<T> auto mat_set, const Form<T, U>& a,
md::mdspan<const scalar_value_t<T>,
md::extents<std::size_t, md::dynamic_extent, 3>>
x,
std::span<const T> constants,
const std::map<std::pair<IntegralType, int>,
std::pair<std::span<const T>, int>>& coefficients,
std::span<const std::int8_t> bc0, std::span<const std::int8_t> bc1)
{
// Integration domain mesh
std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
assert(mesh);
// Test function mesh
auto mesh0 = a.function_spaces().at(0)->mesh();
assert(mesh0);
// Trial function mesh
auto mesh1 = a.function_spaces().at(1)->mesh();
assert(mesh1);
// TODO: Mixed topology with exterior and interior facet integrals.
//
// NOTE: Can't just loop over cell types for interior facet integrals
// because we have a kernel per combination of comparable cell types,
// rather than one per cell type. Also, we need the dofmaps for two
// different cell types at the same time.
const int num_cell_types = mesh->topology()->cell_types().size();
for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
{
// Geometry dofmap and data
mdspan2_t x_dofmap = mesh->geometry().dofmap(cell_type_idx);
// Get dofmap data
std::shared_ptr<const fem::DofMap> dofmap0
= a.function_spaces().at(0)->dofmaps(cell_type_idx);
std::shared_ptr<const fem::DofMap> dofmap1
= a.function_spaces().at(1)->dofmaps(cell_type_idx);
assert(dofmap0);
assert(dofmap1);
auto dofs0 = dofmap0->map();
const int bs0 = dofmap0->bs();
auto dofs1 = dofmap1->map();
const int bs1 = dofmap1->bs();
auto element0 = a.function_spaces().at(0)->elements(cell_type_idx);
assert(element0);
auto element1 = a.function_spaces().at(1)->elements(cell_type_idx);
assert(element1);
fem::DofTransformKernel<T> auto P0
= element0->template dof_transformation_fn<T>(doftransform::standard);
fem::DofTransformKernel<T> auto P1T
= element1->template dof_transformation_right_fn<T>(
doftransform::transpose);
std::span<const std::uint32_t> cell_info0;
std::span<const std::uint32_t> cell_info1;
if (element0->needs_dof_transformations()
or element1->needs_dof_transformations()
or a.needs_facet_permutations())
{
mesh0->topology_mutable()->create_entity_permutations();
mesh1->topology_mutable()->create_entity_permutations();
cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info());
cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info());
}
for (int i = 0; i < a.num_integrals(IntegralType::cell, cell_type_idx); ++i)
{
auto fn = a.kernel(IntegralType::cell, i, cell_type_idx);
assert(fn);
std::span cells = a.domain(IntegralType::cell, i, cell_type_idx);
std::span cells0 = a.domain_arg(IntegralType::cell, 0, i, cell_type_idx);
std::span cells1 = a.domain_arg(IntegralType::cell, 1, i, cell_type_idx);
auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
assert(cells.size() * cstride == coeffs.size());
impl::assemble_cells_matrix(
mat_set, x_dofmap, x, cells, {dofs0, bs0, cells0}, P0,
{dofs1, bs1, cells1}, P1T, bc0, bc1, fn,
md::mdspan(coeffs.data(), cells.size(), cstride), constants,
cell_info0, cell_info1);
}
md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
if (a.needs_facet_permutations())
{
mesh::CellType cell_type = mesh->topology()->cell_types()[cell_type_idx];
int num_facets_per_cell
= mesh::cell_num_entities(cell_type, mesh->topology()->dim() - 1);
mesh->topology_mutable()->create_entity_permutations();
const std::vector<std::uint8_t>& p
= mesh->topology()->get_facet_permutations();
facet_perms = md::mdspan(p.data(), p.size() / num_facets_per_cell,
num_facets_per_cell);
}
for (int i = 0;
i < a.num_integrals(IntegralType::interior_facet, cell_type_idx); ++i)
{
if (num_cell_types > 1)
{
throw std::runtime_error("Interior facet integrals with mixed "
"topology aren't supported yet");
}
using mdspanx22_t
= md::mdspan<const std::int32_t,
md::extents<std::size_t, md::dynamic_extent, 2, 2>>;
using mdspanx2x_t
= md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
md::dynamic_extent>>;
auto fn = a.kernel(IntegralType::interior_facet, i, 0);
assert(fn);
auto& [coeffs, cstride]
= coefficients.at({IntegralType::interior_facet, i});
std::span facets = a.domain(IntegralType::interior_facet, i, 0);
std::span facets0 = a.domain_arg(IntegralType::interior_facet, 0, i, 0);
std::span facets1 = a.domain_arg(IntegralType::interior_facet, 1, i, 0);
assert((facets.size() / 4) * 2 * cstride == coeffs.size());
impl::assemble_interior_facets(
mat_set, x_dofmap, x,
mdspanx22_t(facets.data(), facets.size() / 4, 2, 2),
{*dofmap0, bs0,
mdspanx22_t(facets0.data(), facets0.size() / 4, 2, 2)},
P0,
{*dofmap1, bs1,
mdspanx22_t(facets1.data(), facets1.size() / 4, 2, 2)},
P1T, bc0, bc1, fn,
mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride), constants,
cell_info0, cell_info1, facet_perms);
}
for (auto itg_type : {fem::IntegralType::exterior_facet,
fem::IntegralType::vertex, fem::IntegralType::ridge})
{
md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms
= itg_type == fem::IntegralType::exterior_facet
? facet_perms
: md::mdspan<const std::uint8_t,
md::dextents<std::size_t, 2>>{};
for (int i = 0; i < a.num_integrals(itg_type, cell_type_idx); ++i)
{
if (num_cell_types > 1)
{
throw std::runtime_error("Exterior facet integrals with mixed "
"topology aren't supported yet");
}
using mdspanx2_t
= md::mdspan<const std::int32_t,
md::extents<std::size_t, md::dynamic_extent, 2>>;
auto fn = a.kernel(itg_type, i, 0);
assert(fn);
auto& [coeffs, cstride] = coefficients.at({itg_type, i});
std::span e = a.domain(itg_type, i, 0);
mdspanx2_t entities(e.data(), e.size() / 2, 2);
std::span e0 = a.domain_arg(itg_type, 0, i, 0);
mdspanx2_t entities0(e0.data(), e0.size() / 2, 2);
std::span e1 = a.domain_arg(itg_type, 1, i, 0);
mdspanx2_t entities1(e1.data(), e1.size() / 2, 2);
assert((entities.size() / 2) * cstride == coeffs.size());
impl::assemble_entities(
mat_set, x_dofmap, x, entities, {dofs0, bs0, entities0}, P0,
{dofs1, bs1, entities1}, P1T, bc0, bc1, fn,
md::mdspan(coeffs.data(), entities.extent(0), cstride), constants,
cell_info0, cell_info1, perms);
}
}
}
}
} // namespace dolfinx::fem::impl
DofMap.h
Degree-of-freedom map representations and tools.
utils.h
Functions supporting finite element method operations.
dolfinx::fem::sparsitybuild::cells
void cells(la::SparsityPattern &pattern, std::array< std::span< const std::int32_t >, 2 > cells, std::array< std::reference_wrapper< const DofMap >, 2 > dofmaps)
Iterate over cells and insert entries into sparsity pattern.
Definition sparsitybuild.cpp:16
dolfinx::fem::doftransform::transpose
@ transpose
Transpose.
Definition FiniteElement.h:28
dolfinx::fem::doftransform::standard
@ standard
Standard.
Definition FiniteElement.h:27
dolfinx::fem::IntegralType::vertex
@ vertex
Vertex.
Definition Form.h:43
dolfinx::fem::IntegralType::interior_facet
@ interior_facet
Interior facet.
Definition Form.h:42
dolfinx::fem::IntegralType::ridge
@ ridge
Ridge.
Definition Form.h:44
dolfinx::fem::IntegralType::cell
@ cell
Cell.
Definition Form.h:40
dolfinx::fem::IntegralType::exterior_facet
@ exterior_facet
Facet.
Definition Form.h:41
dolfinx::mesh::CellType
CellType
Cell type identifier.
Definition cell_types.h:22
dolfinx::mesh::cell_num_entities
int cell_num_entities(CellType type, int dim)
Definition cell_types.cpp:139