Skip to content

Commit 0b964a3

Browse files
jorgensdjorgensd
authored
Major refactor of C++ assembly (#150)
* Major refactor. Prepare for submeshes * Use variable * More needs transformation unused added back * Fix element in not implemented yet
1 parent e6c33ea commit 0b964a3

File tree

4 files changed

+363
-324
lines changed

4 files changed

+363
-324
lines changed

‎cpp/assemble_matrix.cpp

+73-54
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,8 @@ void modify_mpc_cell(
163163
// Zero out slave entries in element matrix
164164
// Zero slave row
165165
std::ranges::for_each(local_index[0],
166-
[&Ae, ndim1](const auto dof) {
166+
[&Ae, ndim1](const auto dof)
167+
{
167168
std::ranges::fill_n(
168169
std::next(Ae.data_handle(), ndim1 * dof), ndim1,
169170
0.0);
@@ -276,6 +277,8 @@ void assemble_exterior_facets(
276277
std::span<const std::int32_t>,
277278
const std::span<const T>)>& mat_add_values,
278279
const dolfinx::mesh::Mesh<U>& mesh, std::span<const std::int32_t> facets,
280+
std::span<const std::int32_t> facets0,
281+
std::span<const std::int32_t> facets1,
279282
const std::function<void(const std::span<T>&,
280283
const std::span<const std::uint32_t>&,
281284
std::int32_t, int)>& apply_dof_transformation,
@@ -289,7 +292,8 @@ void assemble_exterior_facets(
289292
const std::uint8_t*)>& kernel,
290293
const std::span<const T> coeffs, int cstride,
291294
const std::vector<T>& constants,
292-
const std::span<const std::uint32_t>& cell_info,
295+
const std::span<const std::uint32_t>& cell_info0,
296+
const std::span<const std::uint32_t>& cell_info1,
293297
const std::shared_ptr<const dolfinx_mpc::MultiPointConstraint<T, U>>& mpc0,
294298
const std::shared_ptr<const dolfinx_mpc::MultiPointConstraint<T, U>>& mpc1)
295299
{
@@ -335,7 +339,8 @@ void assemble_exterior_facets(
335339
std::vector<T> scratch_memory(2 * ndim0 * ndim1 + ndim0 + ndim1);
336340
for (std::size_t l = 0; l < facets.size(); l += 2)
337341
{
338-
342+
const std::int32_t cell0 = facets0[l];
343+
const std::int32_t cell1 = facets1[l];
339344
const std::int32_t cell = facets[l];
340345
const int local_facet = facets[l + 1];
341346

@@ -352,12 +357,12 @@ void assemble_exterior_facets(
352357
std::ranges::fill(Aeb, 0);
353358
kernel(Aeb.data(), coeffs.data() + l / 2 * cstride, constants.data(),
354359
coordinate_dofs.data(), &local_facet, nullptr);
355-
apply_dof_transformation(_Ae, cell_info, cell, ndim1);
356-
apply_dof_transformation_to_transpose(_Ae, cell_info, cell, ndim0);
360+
apply_dof_transformation(_Ae, cell_info0, cell0, ndim1);
361+
apply_dof_transformation_to_transpose(_Ae, cell_info1, cell1, ndim0);
357362

358363
// Zero rows/columns for essential bcs
359-
auto dmap0 = dofmap0.cell_dofs(cell);
360-
auto dmap1 = dofmap1.cell_dofs(cell);
364+
auto dmap0 = dofmap0.cell_dofs(cell0);
365+
auto dmap1 = dofmap1.cell_dofs(cell1);
361366
if (!bc0.empty())
362367
{
363368
for (std::uint32_t i = 0; i < num_dofs0; ++i)
@@ -393,11 +398,11 @@ void assemble_exterior_facets(
393398

394399
// Modify local element matrix Ae and insert contributions into master
395400
// locations
396-
if ((cell_to_slaves[0]->num_links(cell) > 0)
397-
|| (cell_to_slaves[1]->num_links(cell) > 0))
401+
if ((cell_to_slaves[0]->num_links(cell0) > 0)
402+
|| (cell_to_slaves[1]->num_links(cell1) > 0))
398403
{
399404
const std::array<std::span<const std::int32_t>, 2> slaves
400-
= {cell_to_slaves[0]->links(cell), cell_to_slaves[1]->links(cell)};
405+
= {cell_to_slaves[0]->links(cell0), cell_to_slaves[1]->links(cell1)};
401406
const std::array<std::span<const std::int32_t>, 2> dofs = {dmap0, dmap1};
402407
modify_mpc_cell<T>(mat_add_values, num_dofs, Ae, dofs, bs, slaves,
403408
masters, coefficients, is_slave, scratch_memory);
@@ -416,6 +421,8 @@ void assemble_cells_impl(
416421
const std::span<const T>)>& mat_add_values,
417422
const dolfinx::mesh::Geometry<U>& geometry,
418423
std::span<const std::int32_t> active_cells,
424+
std::span<const std::int32_t> active_cells0,
425+
std::span<const std::int32_t> active_cells1,
419426
std::function<void(std::span<T>, const std::span<const std::uint32_t>,
420427
const std::int32_t, const int)>
421428
apply_dof_transformation,
@@ -429,7 +436,8 @@ void assemble_cells_impl(
429436
const std::uint8_t*)>& kernel,
430437
const std::span<const T>& coeffs, int cstride,
431438
const std::vector<T>& constants,
432-
const std::span<const std::uint32_t>& cell_info,
439+
const std::span<const std::uint32_t>& cell_info0,
440+
const std::span<const std::uint32_t>& cell_info1,
433441
const std::shared_ptr<const dolfinx_mpc::MultiPointConstraint<T, U>>& mpc0,
434442
const std::shared_ptr<const dolfinx_mpc::MultiPointConstraint<T, U>>& mpc1)
435443
{
@@ -471,9 +479,12 @@ void assemble_cells_impl(
471479
const std::span<T> _Ae(Aeb);
472480
std::vector<T> scratch_memory(2 * ndim0 * ndim1 + ndim0 + ndim1);
473481

474-
for (std::size_t c = 0; c < active_cells.size(); c++)
482+
for (std::size_t index = 0; index < active_cells0.size(); index++)
475483
{
476-
const std::int32_t cell = active_cells[c];
484+
const std::int32_t cell = active_cells[index];
485+
const std::int32_t cell0 = active_cells0[index];
486+
const std::int32_t cell1 = active_cells1[index];
487+
477488
// Get cell coordinates/geometry
478489
auto x_dofs = MDSPAN_IMPL_STANDARD_NAMESPACE::submdspan(
479490
x_dofmap, cell, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent);
@@ -485,14 +496,14 @@ void assemble_cells_impl(
485496

486497
// Tabulate tensor
487498
std::ranges::fill(Aeb, 0);
488-
kernel(Aeb.data(), coeffs.data() + c * cstride, constants.data(),
499+
kernel(Aeb.data(), coeffs.data() + index * cstride, constants.data(),
489500
coordinate_dofs.data(), nullptr, nullptr);
490-
apply_dof_transformation(_Ae, cell_info, cell, ndim1);
491-
apply_dof_transformation_to_transpose(_Ae, cell_info, cell, ndim0);
501+
apply_dof_transformation(_Ae, cell_info0, cell0, ndim1);
502+
apply_dof_transformation_to_transpose(_Ae, cell_info1, cell1, ndim0);
492503

493504
// Zero rows/columns for essential bcs
494-
std::span<const std::int32_t> dofs0 = dofmap0.cell_dofs(cell);
495-
std::span<const std::int32_t> dofs1 = dofmap1.cell_dofs(cell);
505+
std::span<const std::int32_t> dofs0 = dofmap0.cell_dofs(cell0);
506+
std::span<const std::int32_t> dofs1 = dofmap1.cell_dofs(cell1);
496507
if (!bc0.empty())
497508
{
498509
for (std::uint32_t i = 0; i < num_dofs0; ++i)
@@ -517,11 +528,11 @@ void assemble_cells_impl(
517528

518529
// Modify local element matrix Ae and insert contributions into master
519530
// locations
520-
if ((cell_to_slaves[0]->num_links(cell) > 0)
521-
|| (cell_to_slaves[1]->num_links(cell) > 0))
531+
if ((cell_to_slaves[0]->num_links(cell0) > 0)
532+
|| (cell_to_slaves[1]->num_links(cell1) > 0))
522533
{
523534
const std::array<std::span<const std::int32_t>, 2> slaves
524-
= {cell_to_slaves[0]->links(cell), cell_to_slaves[1]->links(cell)};
535+
= {cell_to_slaves[0]->links(cell0), cell_to_slaves[1]->links(cell1)};
525536
const std::array<std::span<const std::int32_t>, 2> dofs = {dofs0, dofs1};
526537
modify_mpc_cell<T>(mat_add_values, num_dofs, Ae, dofs, bs, slaves,
527538
masters, coefficients, is_slave, scratch_memory);
@@ -543,9 +554,18 @@ void assemble_matrix_impl(
543554
const std::shared_ptr<const dolfinx_mpc::MultiPointConstraint<T, U>>& mpc0,
544555
const std::shared_ptr<const dolfinx_mpc::MultiPointConstraint<T, U>>& mpc1)
545556
{
546-
auto mesh = a.mesh();
557+
// Integration domain mesh
558+
std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
547559
assert(mesh);
548560

561+
// Test function mesh
562+
auto mesh0 = a.function_spaces().at(0)->mesh();
563+
assert(mesh0);
564+
565+
// Trial function mesh
566+
auto mesh1 = a.function_spaces().at(1)->mesh();
567+
assert(mesh1);
568+
549569
// Get dofmap data
550570
std::shared_ptr<const dolfinx::fem::DofMap> dofmap0
551571
= a.function_spaces().at(0)->dofmap();
@@ -573,62 +593,61 @@ void assemble_matrix_impl(
573593
= element1->template dof_transformation_right_fn<T>(
574594
dolfinx::fem::doftransform::transpose);
575595

596+
const int num_cell_types = mesh->topology()->cell_types().size();
597+
if (num_cell_types > 1)
598+
throw std::runtime_error("Not implemented for mixed cell types");
599+
576600
const bool needs_transformation_data
577601
= element0->needs_dof_transformations()
578602
or element1->needs_dof_transformations()
579603
or a.needs_facet_permutations();
580-
std::span<const std::uint32_t> cell_info;
604+
std::span<const std::uint32_t> cell_info0;
605+
std::span<const std::uint32_t> cell_info1;
581606
if (needs_transformation_data)
582607
{
583-
mesh->topology_mutable()->create_entity_permutations();
584-
cell_info = std::span(mesh->topology()->get_cell_permutation_info());
608+
mesh0->topology_mutable()->create_entity_permutations();
609+
mesh1->topology_mutable()->create_entity_permutations();
610+
cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info());
611+
cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info());
585612
}
586613
for (int i : a.integral_ids(dolfinx::fem::IntegralType::cell))
587614
{
588-
const auto& fn = a.kernel(dolfinx::fem::IntegralType::cell, i);
615+
const auto& fn = a.kernel(dolfinx::fem::IntegralType::cell, i, 0);
589616
const auto& [coeffs, cstride]
590617
= coefficients.at({dolfinx::fem::IntegralType::cell, i});
591618
std::span<const std::int32_t> active_cells
592-
= a.domain(dolfinx::fem::IntegralType::cell, i);
619+
= a.domain(dolfinx::fem::IntegralType::cell, i, 0);
620+
std::span<const std::int32_t> active_cells0
621+
= a.domain_arg(dolfinx::fem::IntegralType::cell, 0, i, 0);
622+
std::span<const std::int32_t> active_cells1
623+
= a.domain_arg(dolfinx::fem::IntegralType::cell, 1, i, 0);
593624
assemble_cells_impl<T>(
594625
mat_add_block_values, mat_add_values, mesh->geometry(), active_cells,
595-
apply_dof_transformation, *dofmap0,
626+
active_cells0, active_cells1, apply_dof_transformation, *dofmap0,
596627
apply_dof_transformation_to_transpose, *dofmap1, bc0, bc1, fn, coeffs,
597-
cstride, constants, cell_info, mpc0, mpc1);
628+
cstride, constants, cell_info0, cell_info1, mpc0, mpc1);
598629
}
599630

600631
for (int i : a.integral_ids(dolfinx::fem::IntegralType::exterior_facet))
601632
{
602-
const auto& fn = a.kernel(dolfinx::fem::IntegralType::exterior_facet, i);
633+
const auto& fn = a.kernel(dolfinx::fem::IntegralType::exterior_facet, i, 0);
603634
const auto& [coeffs, cstride]
604635
= coefficients.at({dolfinx::fem::IntegralType::exterior_facet, i});
605636
std::span<const std::int32_t> facets
606-
= a.domain(dolfinx::fem::IntegralType::exterior_facet, i);
607-
assemble_exterior_facets<T>(mat_add_block_values, mat_add_values, *mesh,
608-
facets, apply_dof_transformation, *dofmap0,
609-
apply_dof_transformation_to_transpose, *dofmap1,
610-
bc0, bc1, fn, coeffs, cstride, constants,
611-
cell_info, mpc0, mpc1);
637+
= a.domain(dolfinx::fem::IntegralType::exterior_facet, i, 0);
638+
std::span<const std::int32_t> active_facets0
639+
= a.domain_arg(dolfinx::fem::IntegralType::exterior_facet, 0, i, 0);
640+
std::span<const std::int32_t> active_facets1
641+
= a.domain_arg(dolfinx::fem::IntegralType::exterior_facet, 1, i, 0);
642+
assemble_exterior_facets<T>(
643+
mat_add_block_values, mat_add_values, *mesh, facets, active_facets0,
644+
active_facets1, apply_dof_transformation, *dofmap0,
645+
apply_dof_transformation_to_transpose, *dofmap1, bc0, bc1, fn, coeffs,
646+
cstride, constants, cell_info0, cell_info1, mpc0, mpc1);
612647
}
613648

614-
// if (a.num_integrals(dolfinx::fem::IntegralType::interior_facet) > 0)
615-
// {
616-
// throw std::runtime_error("Not implemented yet");
617-
// // const int tdim = mesh->topology()->dim();
618-
// // mesh->topology_mutable().create_connectivity(tdim - 1, tdim);
619-
// // mesh->topology_mutable().create_entity_permutations();
620-
621-
// // std::function<std::uint8_t(std::size_t)> get_perm;
622-
// // if (a.needs_facet_permutations())
623-
// // {
624-
// // mesh->topology_mutable().create_entity_permutations();
625-
// // const std::vector<std::uint8_t>& perms
626-
// // = mesh->topology()->get_facet_permutations();
627-
// // get_perm = [&perms](std::size_t i) { return perms[i]; };
628-
// // }
629-
// // else
630-
// // get_perm = [](std::size_t) { return 0; };
631-
// }
649+
if (a.integral_ids(dolfinx::fem::IntegralType::interior_facet).size() > 0)
650+
throw std::runtime_error("Not implemented yet");
632651
}
633652
//-----------------------------------------------------------------------------
634653
template <typename T, std::floating_point U>

0 commit comments

Comments
 (0)