File: assemble_expression_impl.h

package info (click to toggle)
fenics-dolfinx 1%3A0.10.0.post4-1exp1
  • links: PTS, VCS
  • area: main
  • in suites: experimental
  • size: 6,028 kB
  • sloc: cpp: 36,535; python: 25,391; makefile: 226; sh: 171; xml: 55
file content (189 lines) | stat: -rw-r--r-- 8,017 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
// Copyright (C) 2025 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later

#pragma once

#include "Expression.h"
#include "FunctionSpace.h"
#include "traits.h"
#include "utils.h"
#include <algorithm>
#include <basix/mdspan.hpp>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <memory>
#include <vector>

namespace dolfinx::fem::impl
{
/// @brief Tabulate an Expression at points.
///
/// Executes an Expression kernel over a list of mesh entities.
///
/// @tparam T Scalar type of expression.
/// @tparam U Geometry type.
/// @param[in,out] values Array in which the tabulated expressions are
/// set. It should have shape `(entities.extent(0), Xshape[0],
/// value_size, num_argument_dofs)`. Storage is row-major. Data is set
/// in `values` (not accumulated).
/// @param[in] fn Expression kernel to execute.
/// @param[in] Xshape Shape `(num_points, geometric dimensions)` of
/// points array at which the expression is evaluated by the kernel.
/// @param[in] value_size Value size of the evaluated expression at a
/// point, e.g. 1 for a scalar field, 3 for a vector field in 3D, 4 for
/// a second-order tensor in 2D.
/// @param[in] num_argument_dofs Dimension of an argument function.
/// Greater than 1 when computing a 1-form expression, e.g. can be used
/// to create a matrix that when applied to a degree-of-freedom vector
/// gives the expression values at the evaluation points.
/// @param[in] x_dofmap Geometry degree-of-freedom map.
/// @param[in] x Geometry coordinate of the mesh.
/// @param[in] coeffs Coefficient data that appears in the Expression.
/// Shape is `(num_cells, coeff data per cell)`. Usually packed using
/// fem::pack_coefficients.
/// @param[in] constants Constant (coefficient) data that appears in
/// expression. Usually packed using fem::pack_constants.
/// @param[in] entities Mesh entities to evaluate the expression over.
/// For expressions executed on cells, rank is 1 and size is the number
/// of cells. For expressions executed on facets rank is 2, and shape is
/// `(num_facets, 2)`, where `entities[i, 0]` is the cell index and
/// `entities[i, 1]` is the local index of the facet relative to the
/// cell.
/// @param[in] cell_info Cell orientation data for use in `P0`.
/// @param[in] P0 Degree-of-freedom transformation function. Applied when
/// expressions includes an argument function that requires a
/// transformation.
template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(
    std::span<T> values, fem::FEkernel<T> auto fn,
    std::array<std::size_t, 2> Xshape, std::size_t value_size,
    std::size_t num_argument_dofs,
    md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> x_dofmap,
    std::span<const scalar_value_t<T>> x,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    std::span<const T> constants, fem::MDSpan2 auto entities,
    std::span<const std::uint32_t> cell_info,
    fem::DofTransformKernel<T> auto P0)
{
  static_assert(entities.rank() == 1 or entities.rank() == 2);

  // Create data structures used in evaluation
  std::vector<U> coord_dofs(3 * x_dofmap.extent(1));

  // Iterate over cells and 'assemble' into values
  int size0 = Xshape[0] * value_size;
  std::vector<T> values_local(size0 * num_argument_dofs, 0);
  std::size_t offset = values_local.size();
  for (std::size_t e = 0; e < entities.extent(0); ++e)
  {
    std::ranges::fill(values_local, 0);
    if constexpr (entities.rank() == 1)
    {
      std::int32_t entity = entities(e);
      auto x_dofs = md::submdspan(x_dofmap, entity, md::full_extent);
      for (std::size_t i = 0; i < x_dofs.size(); ++i)
      {
        std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                    std::next(coord_dofs.begin(), 3 * i));
      }
      fn(values_local.data(), &coeffs(e, 0), constants.data(),
         coord_dofs.data(), nullptr, nullptr, nullptr);
    }
    else
    {
      std::int32_t entity = entities(e, 0);
      auto x_dofs = md::submdspan(x_dofmap, entity, md::full_extent);
      for (std::size_t i = 0; i < x_dofs.size(); ++i)
      {
        std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                    std::next(coord_dofs.begin(), 3 * i));
      }
      fn(values_local.data(), &coeffs(e, 0), constants.data(),
         coord_dofs.data(), &entities(e, 1), nullptr, nullptr);
    }

    P0(values_local, cell_info, e, size0);
    for (std::size_t j = 0; j < values_local.size(); ++j)
      values[e * offset + j] = values_local[j];
  }
}

/// @brief Tabulate an Expression at points.
///
/// Executes an Expression kernel over a list of mesh entities.
///
/// @tparam T Scalar type
/// @tparam U Geometry type
/// @param[in,out] values Array in which the tabulated expressions are
/// set. If `V` is available, it should have shape (`entities.extent(0),
/// Xshape[0], value_size, dim(V))`. Otherwise is has shape
/// (`entities.extent(0), Xshape[0], value_size)`. Storage is row-major.
/// Data is set (not accumulated).
/// @param[in] fn Expression kernel to execute.
/// @param[in] Xshape Shape `(num_points, geometric dimensions)` of
/// points array at which the expression is evaluated by the kernel.
/// @param[in] value_size Value size of the evaluated expression at a
/// point, e.g. 1 for a scalar field and 3 for a vector field in 3D.
/// @param[in] coeffs Coefficient data that appears in the Expression.
/// Shape is `(num_cells, coeff data per cell)`. Usually packed using
/// fem::pack_coefficients.
/// @param[in] constants Constant (coefficient) data that appears in
/// expression. Usually packed using fem::pack_constants.
/// @param[in] mesh Mesh to execute the expression kernel on.
/// @param[in] entities Mesh entities to evaluate the expression over.
/// For expressions executed on cells, rank is 1 and size is the number
/// of cells. For expressions executed on facets rank is 2, and shape is
/// `(num_facets, 2)`, where `entities[i, 0]` is the cell index and
/// `entities[i, 1]` is the local index of the facet relative to the
/// cell.
/// @param[in] element Argument element and argument space dimension.
/// Used to computed a 1-form expression, e.g. can be used to create a
/// matrix that when applied to a degree-of-freedom vector gives the
/// expression values at the evaluation points.
template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(
    std::span<T> values, fem::FEkernel<T> auto fn,
    std::array<std::size_t, 2> Xshape, std::size_t value_size,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    std::span<const T> constants, const mesh::Mesh<U>& mesh,
    fem::MDSpan2 auto entities,
    std::optional<
        std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
        element)
{
  std::function<void(std::span<T>, std::span<const std::uint32_t>, std::int32_t,
                     int)>
      post_dof_transform
      = [](std::span<T>, std::span<const std::uint32_t>, std::int32_t, int)
  {
    // Do nothing
  };

  std::shared_ptr<const mesh::Topology> topology = mesh.topology();
  assert(topology);
  std::size_t num_argument_dofs = 1;
  std::span<const std::uint32_t> cell_info;
  if (element)
  {
    num_argument_dofs = element->second;
    if (element->first.get().needs_dof_transformations())
    {
      mesh.topology_mutable()->create_entity_permutations();
      cell_info = std::span(topology->get_cell_permutation_info());
      post_dof_transform
          = element->first.get().template dof_transformation_right_fn<T>(
              doftransform::transpose);
    }
  }

  tabulate_expression<T, U>(values, fn, Xshape, value_size, num_argument_dofs,
                            mesh.geometry().dofmap(), mesh.geometry().x(),
                            coeffs, constants, entities, cell_info,
                            post_dof_transform);
}
} // namespace dolfinx::fem::impl