From bf641cbd267e277a12e7c490b8dd470ec32de032 Mon Sep 17 00:00:00 2001
From: Katrin Mang <mang@ifam.uni-hannover.de>
Date: Fri, 14 Feb 2020 16:46:28 +0100
Subject: [PATCH] add split_by_block into index_set

---
 doc/news/changes/minor/20200214KatrinMang |  4 ++
 include/deal.II/base/index_set.h          |  8 +++
 source/base/index_set.cc                  | 18 ++++++
 tests/base/index_set_31.cc                | 71 +++++++++++++++++++++++
 tests/base/index_set_31.output            |  7 +++
 5 files changed, 108 insertions(+)
 create mode 100644 doc/news/changes/minor/20200214KatrinMang
 create mode 100644 tests/base/index_set_31.cc
 create mode 100644 tests/base/index_set_31.output

diff --git a/doc/news/changes/minor/20200214KatrinMang b/doc/news/changes/minor/20200214KatrinMang
new file mode 100644
index 0000000000..4966c6c58a
--- /dev/null
+++ b/doc/news/changes/minor/20200214KatrinMang
@@ -0,0 +1,4 @@
+New: An IndexSet::split_by_block() function is added which partitions the set indices for example based on a block structure into blocks.
+IndexSet should be defined with respect to the global block structure. In particular, dofs_per_block should be defined consistently across MPI ranks.
+<br>
+(Katrin Mang, 2020/02/14)
diff --git a/include/deal.II/base/index_set.h b/include/deal.II/base/index_set.h
index b5ada5fb26..ca00bf3288 100644
--- a/include/deal.II/base/index_set.h
+++ b/include/deal.II/base/index_set.h
@@ -344,6 +344,14 @@ public:
   IndexSet
   get_view(const size_type begin, const size_type end) const;
 
+  /**
+   * Split the set indices represented by this object into blocks
+   * given by the @p dofs_per_block structure.
+   */
+  std::vector<IndexSet>
+  split_by_block(
+    const std::vector<types::global_dof_index> &dofs_per_block) const;
+
   /**
    * Remove all elements contained in @p other from this set. In other words,
    * if $x$ is the current object and $o$ the argument, then we compute $x
diff --git a/source/base/index_set.cc b/source/base/index_set.cc
index 6f986b956f..67892e6482 100644
--- a/source/base/index_set.cc
+++ b/source/base/index_set.cc
@@ -232,7 +232,25 @@ IndexSet::get_view(const size_type begin, const size_type end) const
   return result;
 }
 
+std::vector<IndexSet>
+IndexSet::split_by_block(
+  const std::vector<types::global_dof_index> &dofs_per_block) const
+{
+  std::vector<IndexSet> partitioned;
+  const unsigned int    n_block_dofs = dofs_per_block.size();
 
+  partitioned.reserve(n_block_dofs);
+  types::global_dof_index start = 0;
+  types::global_dof_index sum   = 0;
+  for (const auto n_block_dofs : dofs_per_block)
+    {
+      partitioned.push_back(this->get_view(start, start + n_block_dofs));
+      start += n_block_dofs;
+      sum += partitioned.back().size();
+    }
+  AssertDimension(sum, this->size());
+  return partitioned;
+}
 
 void
 IndexSet::subtract_set(const IndexSet &other)
diff --git a/tests/base/index_set_31.cc b/tests/base/index_set_31.cc
new file mode 100644
index 0000000000..2089371c8d
--- /dev/null
+++ b/tests/base/index_set_31.cc
@@ -0,0 +1,71 @@
+// ---------------------------------------------------------------------
+//
+// Copyright (C) 2020 by the deal.II authors
+//
+// This file is part of the deal.II library.
+//
+// The deal.II library is free software; you can use it, redistribute
+// it, and/or modify it under the terms of the GNU Lesser General
+// Public License as published by the Free Software Foundation; either
+// version 2.1 of the License, or (at your option) any later version.
+// The full text of the license can be found in the file LICENSE.md at
+// the top level directory of deal.II.
+//
+// ---------------------------------------------------------------------
+
+
+// test IndexSet::split_by_block ()
+
+#include <deal.II/base/index_set.h>
+
+#include <deal.II/dofs/dof_handler.h>
+#include <deal.II/dofs/dof_renumbering.h>
+#include <deal.II/dofs/dof_tools.h>
+
+#include <deal.II/fe/fe_q.h>
+#include <deal.II/fe/fe_system.h>
+
+#include <deal.II/grid/grid_generator.h>
+#include <deal.II/grid/tria.h>
+
+#include "../tests.h"
+
+template <int dim>
+void
+test()
+{
+  Triangulation<dim> triangulation;
+  GridGenerator::hyper_cube(triangulation);
+  triangulation.refine_global(1);
+
+  FESystem<dim>   fe(FE_Q<dim>(2), dim, FE_Q<dim>(1), 1);
+  DoFHandler<dim> dof_handler(triangulation);
+  dof_handler.distribute_dofs(fe);
+
+  std::vector<unsigned int> block_component(dim + 1, 0);
+  block_component[dim] = 1;
+  DoFRenumbering::component_wise(dof_handler, block_component);
+
+  const std::vector<types::global_dof_index> dofs_per_block =
+    DoFTools::count_dofs_per_fe_block(dof_handler, block_component);
+
+  const std::vector<IndexSet> partition =
+    dof_handler.locally_owned_dofs().split_by_block(dofs_per_block);
+  for (unsigned int i = 0; i < 2; ++i)
+    {
+      partition[i].print(deallog);
+    }
+
+  deallog << "OK" << std::endl;
+}
+
+
+
+int
+main()
+{
+  initlog();
+
+  test<2>();
+  test<3>();
+}
diff --git a/tests/base/index_set_31.output b/tests/base/index_set_31.output
new file mode 100644
index 0000000000..f379192b08
--- /dev/null
+++ b/tests/base/index_set_31.output
@@ -0,0 +1,7 @@
+
+DEAL::{[0,49]}
+DEAL::{[0,8]}
+DEAL::OK
+DEAL::{[0,374]}
+DEAL::{[0,26]}
+DEAL::OK
-- 
2.39.5