From: Peter Munch <peterrmuench@gmail.com>
Date: Tue, 18 Jan 2022 11:22:42 +0000 (+0100)
Subject: FEEval: allow to reinit with a batch with arbitrary cells
X-Git-Tag: v9.4.0-rc1~545^2
X-Git-Url: https://gitweb.dealii.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=refs%2Fpull%2F13260%2Fhead;p=dealii.git

FEEval: allow to reinit with a batch with arbitrary cells
---

diff --git a/include/deal.II/lac/la_vector.templates.h b/include/deal.II/lac/la_vector.templates.h
index c8713aea5b..5d68773869 100644
--- a/include/deal.II/lac/la_vector.templates.h
+++ b/include/deal.II/lac/la_vector.templates.h
@@ -71,11 +71,7 @@ namespace LinearAlgebra
     Assert(dynamic_cast<const Vector<Number> *>(&V) != nullptr,
            ExcVectorTypeNotCompatible());
 
-    // Downcast V. If fails, throws an exception.
     const Vector<Number> &down_V = dynamic_cast<const Vector<Number> &>(V);
-    Assert(down_V.size() == this->size(),
-           ExcMessage(
-             "Cannot add two vectors with different numbers of elements"));
 
     ReadWriteVector<Number>::reinit(down_V, omit_zeroing_entries);
   }
diff --git a/include/deal.II/matrix_free/fe_evaluation.h b/include/deal.II/matrix_free/fe_evaluation.h
index 6d9b4290a7..a5b7846cde 100644
--- a/include/deal.II/matrix_free/fe_evaluation.h
+++ b/include/deal.II/matrix_free/fe_evaluation.h
@@ -778,6 +778,33 @@ protected:
    * MatrixFree object was given at initialization.
    */
   mutable std::vector<types::global_dof_index> local_dof_indices;
+
+  /**
+   * A temporary data structure for the Jacobian information necessary if
+   * the reinit() function is used that allows to construct user batches.
+   */
+  AlignedVector<Tensor<2, dim, VectorizedArrayType>> jacobian_data;
+
+  /**
+   * A temporary data structure for the Jacobian determinant necessary if
+   * the reinit() function is used that allows to construct user batches.
+   */
+  AlignedVector<VectorizedArrayType> J_value_data;
+
+  /**
+   * A temporary data structure for the gradients of the inverse Jacobian
+   * transformation necessary if the reinit() function is used that allows
+   * to construct user batches.
+   */
+  AlignedVector<
+    Tensor<1, dim *(dim + 1) / 2, Tensor<1, dim, VectorizedArrayType>>>
+    jacobian_gradients_data;
+
+  /**
+   * A temporary data structure for the quadrature-point locations necessary if
+   * the reinit() function is used that allows to construct user batches.
+   */
+  AlignedVector<Point<dim, VectorizedArrayType>> quadrature_points_data;
 };
 
 
@@ -2150,6 +2177,15 @@ public:
   void
   reinit(const unsigned int cell_batch_index);
 
+  /**
+   * Similar as the above function but allowing to define customized cell
+   * batches on the fly. A cell batch is defined by the (matrix-free) index of
+   * its cells: see also the documentation of get_cell_ids () or
+   * get_cell_or_face_ids ().
+   */
+  void
+  reinit(const std::array<unsigned int, VectorizedArrayType::size()> &cell_ids);
+
   /**
    * Initialize the data to the current cell using a TriaIterator object as
    * usual in FEValues. The argument is either of type
@@ -3485,7 +3521,8 @@ FEEvaluationBase<dim, n_components_, Number, is_face, VectorizedArrayType>::
     values_dofs[c] = const_cast<VectorizedArrayType *>(this->values_dofs) +
                      c * dofs_per_component;
 
-  if (this->dof_info->index_storage_variants
+  if (this->cell != numbers::invalid_unsigned_int &&
+      this->dof_info->index_storage_variants
           [is_face ? this->dof_access_index :
                      internal::MatrixFreeFunctions::DoFInfo::dof_access_cell]
           [this->cell] == internal::MatrixFreeFunctions::DoFInfo::
@@ -6904,6 +6941,224 @@ FEEvaluation<dim,
 
 
 
+template <int dim,
+          int fe_degree,
+          int n_q_points_1d,
+          int n_components_,
+          typename Number,
+          typename VectorizedArrayType>
+inline void
+FEEvaluation<dim,
+             fe_degree,
+             n_q_points_1d,
+             n_components_,
+             Number,
+             VectorizedArrayType>::
+  reinit(const std::array<unsigned int, VectorizedArrayType::size()> &cell_ids)
+{
+  Assert(this->mapped_geometry == nullptr,
+         ExcMessage("FEEvaluation was initialized without a matrix-free object."
+                    " Integer indexing is not possible"));
+  if (this->mapped_geometry != nullptr)
+    return;
+
+  Assert(this->dof_info != nullptr, ExcNotInitialized());
+  Assert(this->mapping_data != nullptr, ExcNotInitialized());
+
+  this->cell     = numbers::invalid_unsigned_int;
+  this->cell_ids = cell_ids;
+
+  // determine type of cell batch
+  this->cell_type = internal::MatrixFreeFunctions::GeometryType::cartesian;
+
+  for (unsigned int v = 0; v < VectorizedArrayType::size(); ++v)
+    {
+      const unsigned int cell_index = cell_ids[v];
+
+      if (cell_index == numbers::invalid_unsigned_int)
+        continue;
+
+      this->cell_type =
+        std::max(this->cell_type,
+                 this->matrix_free->get_mapping_info().get_cell_type(
+                   cell_index / VectorizedArrayType::size()));
+    }
+
+  // allocate memory for internal data storage
+
+  if (this->cell_type <= internal::MatrixFreeFunctions::GeometryType::affine)
+    {
+      if (this->mapping_data->jacobians[0].size() > 0)
+        this->jacobian_data.resize_fast(2);
+
+      if (this->mapping_data->JxW_values.size() > 0)
+        this->J_value_data.resize_fast(1);
+
+      if (this->mapping_data->jacobian_gradients[0].size() > 0)
+        this->jacobian_gradients_data.resize_fast(1);
+
+      if (this->mapping_data->quadrature_points.size() > 0)
+        this->quadrature_points_data.resize_fast(1);
+    }
+  else
+    {
+      if (this->mapping_data->jacobians[0].size() > 0)
+        this->jacobian_data.resize_fast(this->n_quadrature_points);
+
+      if (this->mapping_data->JxW_values.size() > 0)
+        this->J_value_data.resize_fast(this->n_quadrature_points);
+
+      if (this->mapping_data->jacobian_gradients[0].size() > 0)
+        this->jacobian_gradients_data.resize_fast(this->n_quadrature_points);
+
+      if (this->mapping_data->quadrature_points.size() > 0)
+        this->quadrature_points_data.resize_fast(this->n_quadrature_points);
+    }
+
+  // set pointers to internal data storage
+  this->jacobian           = this->jacobian_data.data();
+  this->J_value            = this->J_value_data.data();
+  this->jacobian_gradients = this->jacobian_gradients_data.data();
+  this->quadrature_points  = this->quadrature_points_data.data();
+
+  // fill internal data storage lane by lane
+  for (unsigned int v = 0; v < VectorizedArrayType::size(); ++v)
+    {
+      const unsigned int cell_index = cell_ids[v];
+
+      if (cell_index == numbers::invalid_unsigned_int)
+        continue;
+
+      const unsigned int cell_batch_index =
+        cell_index / VectorizedArrayType::size();
+      const unsigned int offsets =
+        this->mapping_data->data_index_offsets[cell_batch_index];
+      const unsigned int lane = cell_index % VectorizedArrayType::size();
+
+      if (this->cell_type <=
+          internal::MatrixFreeFunctions::GeometryType::affine)
+        {
+          // case that all cells are Cartesian or affine
+          const unsigned int q = 0;
+
+          if (this->mapping_data->JxW_values.size() > 0)
+            this->J_value_data[q][v] =
+              this->mapping_data->JxW_values[offsets + q][lane];
+
+          if (this->mapping_data->jacobians[0].size() > 0)
+            for (unsigned int q = 0; q < 2; ++q)
+              for (unsigned int i = 0; i < dim; ++i)
+                for (unsigned int j = 0; j < dim; ++j)
+                  this->jacobian_data[q][i][j][v] =
+                    this->mapping_data->jacobians[0][offsets + q][i][j][lane];
+
+          if (this->mapping_data->jacobian_gradients[0].size() > 0)
+            for (unsigned int i = 0; i < dim * (dim + 1) / 2; ++i)
+              for (unsigned int j = 0; j < dim; ++j)
+                this->jacobian_gradients_data[q][i][j][v] =
+                  this->mapping_data
+                    ->jacobian_gradients[0][offsets + q][i][j][lane];
+
+          if (this->mapping_data->quadrature_points.size() > 0)
+            for (unsigned int i = 0; i < dim; ++i)
+              this->quadrature_points_data[q][i][v] =
+                this->mapping_data->quadrature_points
+                  [this->mapping_data
+                     ->quadrature_point_offsets[cell_batch_index] +
+                   q][i][lane];
+        }
+      else
+        {
+          // general case that at least one cell is not Cartesian or affine
+          const auto cell_type =
+            this->matrix_free->get_mapping_info().get_cell_type(
+              cell_batch_index);
+
+          for (unsigned int q = 0; q < this->n_quadrature_points; ++q)
+            {
+              const unsigned int q_src =
+                (cell_type <=
+                 internal::MatrixFreeFunctions::GeometryType::affine) ?
+                  0 :
+                  q;
+
+              if (this->mapping_data->JxW_values.size() > 0)
+                this->J_value_data[q][v] =
+                  this->mapping_data->JxW_values[offsets + q_src][lane];
+
+              if (this->mapping_data->jacobians[0].size() > 0)
+                for (unsigned int i = 0; i < dim; ++i)
+                  for (unsigned int j = 0; j < dim; ++j)
+                    this->jacobian_data[q][i][j][v] =
+                      this->mapping_data
+                        ->jacobians[0][offsets + q_src][i][j][lane];
+
+              if (this->mapping_data->jacobian_gradients[0].size() > 0)
+                for (unsigned int i = 0; i < dim * (dim + 1) / 2; ++i)
+                  for (unsigned int j = 0; j < dim; ++j)
+                    this->jacobian_gradients_data[q][i][j][v] =
+                      this->mapping_data
+                        ->jacobian_gradients[0][offsets + q_src][i][j][lane];
+
+              if (this->mapping_data->quadrature_points.size() > 0)
+                {
+                  if (cell_type <=
+                      internal::MatrixFreeFunctions::GeometryType::affine)
+                    {
+                      // affine case: quadrature points are not available but
+                      // have to be computed from the the corner point and the
+                      // Jacobian
+                      Point<dim, VectorizedArrayType> point =
+                        this->mapping_data->quadrature_points
+                          [this->mapping_data
+                             ->quadrature_point_offsets[cell_batch_index] +
+                           0];
+
+                      const Tensor<2, dim, VectorizedArrayType> &jac =
+                        this->mapping_data->jacobians[0][offsets + 1];
+                      if (cell_type == internal::MatrixFreeFunctions::cartesian)
+                        for (unsigned int d = 0; d < dim; ++d)
+                          point[d] +=
+                            jac[d][d] *
+                            static_cast<Number>(
+                              this->descriptor->quadrature.point(q)[d]);
+                      else
+                        for (unsigned int d = 0; d < dim; ++d)
+                          for (unsigned int e = 0; e < dim; ++e)
+                            point[d] +=
+                              jac[d][e] *
+                              static_cast<Number>(
+                                this->descriptor->quadrature.point(q)[e]);
+
+                      for (unsigned int i = 0; i < dim; ++i)
+                        this->quadrature_points_data[q][i][v] = point[i][lane];
+                    }
+                  else
+                    {
+                      // general case: quadrature points are available
+                      for (unsigned int i = 0; i < dim; ++i)
+                        this->quadrature_points_data[q][i][v] =
+                          this->mapping_data->quadrature_points
+                            [this->mapping_data
+                               ->quadrature_point_offsets[cell_batch_index] +
+                             q][i][lane];
+                    }
+                }
+            }
+        }
+    }
+
+#  ifdef DEBUG
+  this->is_reinitialized           = true;
+  this->dof_values_initialized     = false;
+  this->values_quad_initialized    = false;
+  this->gradients_quad_initialized = false;
+  this->hessians_quad_initialized  = false;
+#  endif
+}
+
+
+
 template <int dim,
           int fe_degree,
           int n_q_points_1d,
diff --git a/tests/matrix_free/fe_evaluation_renumbered.cc b/tests/matrix_free/fe_evaluation_renumbered.cc
new file mode 100644
index 0000000000..409b629125
--- /dev/null
+++ b/tests/matrix_free/fe_evaluation_renumbered.cc
@@ -0,0 +1,231 @@
+// ---------------------------------------------------------------------
+//
+// Copyright (C) 2021 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.
+//
+// ---------------------------------------------------------------------
+
+#include <deal.II/base/quadrature_lib.h>
+
+#include <deal.II/dofs/dof_handler.h>
+
+#include <deal.II/fe/fe_dgq.h>
+#include <deal.II/fe/mapping_q.h>
+
+#include <deal.II/grid/grid_generator.h>
+
+#include <deal.II/matrix_free/fe_evaluation.h>
+#include <deal.II/matrix_free/matrix_free.h>
+
+#include <deal.II/multigrid/mg_constrained_dofs.h>
+
+#include <deal.II/numerics/vector_tools.h>
+
+#include "../tests.h"
+
+
+// Test FEEvaluation::reinit() that thaes std::array of indices.
+
+
+
+template <int dim>
+class RightHandSideFunction : public Function<dim>
+{
+public:
+  RightHandSideFunction(const unsigned int n_components = 1)
+    : Function<dim>(n_components)
+  {}
+
+  virtual double
+  value(const Point<dim> &p, const unsigned int component = 0) const
+  {
+    return p[component % dim] * p[component % dim];
+  }
+};
+
+template <int dim, int fe_degree, int n_points, typename Number>
+void
+test(const unsigned int n_refinements, const unsigned int geometry_type)
+{
+  using VectorizedArrayType = VectorizedArray<Number>;
+
+  using VectorType = LinearAlgebra::Vector<Number>;
+
+  Triangulation<dim> tria;
+
+  if (geometry_type == 0)
+    GridGenerator::hyper_cube(tria);
+  else if (geometry_type == 1)
+    GridGenerator::hyper_ball(tria);
+  else
+    AssertThrow(false, ExcNotImplemented());
+
+  tria.refine_global(n_refinements);
+
+  FE_Q<dim>       fe(fe_degree);
+  DoFHandler<dim> dof_handler(tria);
+  dof_handler.distribute_dofs(fe);
+
+  MappingQ<dim> mapping(1);
+  QGauss<1>     quad(n_points);
+
+  AffineConstraints<Number> constraint;
+
+  using MF = MatrixFree<dim, Number, VectorizedArrayType>;
+
+  typename MF::AdditionalData additional_data;
+  additional_data.mapping_update_flags =
+    update_values | update_quadrature_points;
+  additional_data.tasks_parallel_scheme = MF::AdditionalData::none;
+
+  MF matrix_free;
+  matrix_free.reinit(mapping, dof_handler, constraint, quad, additional_data);
+
+  VectorType src, dst;
+
+  matrix_free.initialize_dof_vector(src);
+  matrix_free.initialize_dof_vector(dst);
+
+  VectorTools::interpolate(mapping,
+                           dof_handler,
+                           RightHandSideFunction<dim>(1),
+                           src);
+
+  std::vector<std::tuple<unsigned int, unsigned int, unsigned int>> indices;
+
+  for (unsigned int cell = 0; cell < matrix_free.n_cell_batches(); ++cell)
+    {
+      for (unsigned int i = 0;
+           i < matrix_free.n_active_entries_per_cell_batch(cell);
+           ++i)
+        indices.emplace_back(indices.size(), cell, i);
+    }
+
+  // if(false)
+  std::sort(indices.begin(), indices.end(), [](const auto &a, const auto &b) {
+    if (std::get<2>(a) != std::get<2>(b))
+      return std::get<2>(a) < std::get<2>(b);
+
+    return std::get<1>(a) < std::get<1>(b);
+  });
+
+  // collect reference data with normal reinit function
+
+  // quadrature points -> mapping
+  std::vector<std::vector<Point<dim>>> quadrature_points_ref;
+
+  // dof values -> read/write operation
+  std::vector<std::vector<Number>> values_ref;
+
+  matrix_free.template cell_loop<VectorType, VectorType>(
+    [&](const auto &matrix_free, auto &dst, const auto &src, auto range) {
+      FEEvaluation<dim, fe_degree, n_points, 1, Number, VectorizedArrayType>
+        phi(matrix_free);
+
+      for (unsigned int cell = range.first; cell < range.second; ++cell)
+        {
+          phi.reinit(cell);
+          phi.read_dof_values(src);
+
+          for (unsigned int v = 0;
+               v < matrix_free.n_active_entries_per_cell_batch(cell);
+               ++v)
+            {
+              std::vector<Point<dim>> points;
+
+              for (unsigned int i = 0; i < phi.n_q_points; ++i)
+                {
+                  auto temp_v = phi.quadrature_point(i);
+
+                  Point<dim> temp;
+                  for (unsigned int d = 0; d < dim; ++d)
+                    temp[d] = temp_v[d][v];
+
+                  points.emplace_back(temp);
+                }
+
+              quadrature_points_ref.emplace_back(points);
+
+              std::vector<Number> values;
+
+              for (unsigned int i = 0; i < phi.dofs_per_cell; ++i)
+                values.emplace_back(phi.get_dof_value(i)[v]);
+              values_ref.emplace_back(values);
+            }
+        }
+    },
+    dst,
+    src);
+
+  // test variable reinit() function
+  {
+    FEEvaluation<dim, fe_degree, n_points, 1, Number, VectorizedArrayType> phi(
+      matrix_free);
+
+    for (unsigned int v = 0; v < indices.size();
+         v += VectorizedArrayType::size())
+      {
+        std::array<unsigned int, VectorizedArrayType::size()> indices_;
+
+        indices_.fill(numbers::invalid_unsigned_int);
+
+        const unsigned int n_lanes_filled =
+          std::min(v + VectorizedArrayType::size(), indices.size()) - v;
+
+        for (unsigned int i = v, c = 0; i < v + n_lanes_filled; ++i, ++c)
+          indices_[c] = std::get<1>(indices[i]) * VectorizedArrayType::size() +
+                        std::get<2>(indices[i]);
+
+        phi.reinit(indices_);
+        phi.read_dof_values(src);
+
+        for (unsigned int i = v, c = 0; i < v + n_lanes_filled; ++i, ++c)
+          {
+            std::vector<Point<dim>> points;
+
+            for (unsigned int i = 0; i < phi.n_q_points; ++i)
+              {
+                auto temp_v = phi.quadrature_point(i);
+
+                Point<dim> temp;
+                for (unsigned int d = 0; d < dim; ++d)
+                  temp[d] = temp_v[d][c];
+
+                points.emplace_back(temp);
+              }
+
+            // perform comparison
+
+            Assert(points == quadrature_points_ref[std::get<0>(indices[i])],
+                   ExcInternalError());
+
+            std::vector<Number> values;
+
+            for (unsigned int i = 0; i < phi.dofs_per_cell; ++i)
+              values.emplace_back(phi.get_dof_value(i)[c]);
+
+            // perform comparison
+            Assert(values == values_ref[std::get<0>(indices[i])],
+                   ExcInternalError());
+          }
+      }
+  }
+
+  deallog << "OK!" << std::endl;
+}
+
+int
+main()
+{
+  initlog();
+  test<2, 1, 2, double>(3, 0);
+  test<2, 1, 2, double>(3, 1);
+}
diff --git a/tests/matrix_free/fe_evaluation_renumbered.output b/tests/matrix_free/fe_evaluation_renumbered.output
new file mode 100644
index 0000000000..043986ad9a
--- /dev/null
+++ b/tests/matrix_free/fe_evaluation_renumbered.output
@@ -0,0 +1,3 @@
+
+DEAL::OK!
+DEAL::OK!