Add write_meshtags

cpp/demo/checkpointing/main.cpp

@@ -30,6 +30,45 @@ int main(int argc, char* argv[])
       MPI_COMM_WORLD, {{{0.0, 0.0}, {1.0, 1.0}}}, {4, 4},
       mesh::CellType::quadrilateral, part));
 
+  // Create cell meshtags
+  auto geometry = mesh->geometry();
+  auto topology = mesh->topology();
+
+  int dim = geometry.dim();
+  topology->create_entities(dim);
+  const std::shared_ptr<const dolfinx::common::IndexMap> topo_imap
+      = topology->index_map(dim);
+
+  std::uint32_t num_entities = topo_imap->size_local();
+
+  auto cmap = mesh->geometry().cmap();
+  auto geom_layout = cmap.create_dof_layout();
+  std::uint32_t num_dofs_per_entity = geom_layout.num_entity_closure_dofs(dim);
+
+  std::vector<int32_t> entities_array(num_entities * num_dofs_per_entity);
+  std::vector<int32_t> entities_offsets(num_entities + 1);
+  std::uint64_t offset = topo_imap->local_range()[0];
+  std::vector<std::uint64_t> values(num_entities);
+
+  for (std::uint32_t i = 0; i < num_entities; ++i)
+  {
+    values[i] = i + offset;
+  }
+
+  auto entities = topology->connectivity(dim, 0);
+
+  for (int i = 0; i < (int)num_entities + 1; ++i)
+    entities_offsets[i] = entities->offsets()[i];
+
+  for (int i = 0; i < (int)(num_entities * num_dofs_per_entity); ++i)
+    entities_array[i] = entities->array()[i];
+
+  graph::AdjacencyList<std::int32_t> entities_local(entities_array,
+                                                    entities_offsets);
+
+  auto meshtags = std::make_shared<mesh::MeshTags<std::uint64_t>>(
+      mesh::create_meshtags<std::uint64_t>(topology, dim, entities_local, values));
+
   try
   {
     // Set up ADIOS2 IO and Engine
@@ -40,6 +79,7 @@ int main(int argc, char* argv[])
     adios2::Engine engine = io.Open("mesh.bp", adios2::Mode::Append);
 
     io::native::write_mesh(io, engine, *mesh);
+    io::native::write_meshtags<float, std::uint64_t>(io, engine, *mesh, *meshtags);
 
     engine.Close();
   }

cpp/dolfinx/io/checkpointing.h

@@ -12,12 +12,46 @@
 #include <adios2.h>
 #include <basix/finite-element.h>
 #include <dolfinx/mesh/Mesh.h>
+#include <dolfinx/mesh/MeshTags.h>
 #include <dolfinx/mesh/utils.h>
 #include <mpi.h>
 
 /// @file checkpointing.h
 /// @brief ADIOS2 based checkpointing
 
+namespace
+{
+/// @brief Write a particular attribute incrementally.
+/// For example, to write a new meshtag, fetch the name attribute if it exists
+/// and append the name, otherwise create the attribute.
+/// @tparam T ADIOS2 supported type
+/// @param io ADIOS2 IO object
+/// @param name Name of the attribute
+/// @param value Value of the attribute to write
+/// @param var_name Variable to which this attribute is associated with
+/// @return Return the IO attribute
+template <class T>
+adios2::Attribute<T> define_attr(adios2::IO& io, const std::string& name,
+                                 T& value, std::string var_name = "")
+{
+  bool modifiable = true;
+  if (adios2::Attribute<T> attr = io.InquireAttribute<T>(name); attr)
+  {
+    std::vector<T> data = attr.Data();
+    data.push_back(value);
+    return io.DefineAttribute<T>(name, data.data(), data.size(), var_name, "/",
+                                 modifiable);
+  }
+  else
+  {
+    std::vector<T> data{value};
+    return io.DefineAttribute<T>(name, data.data(), data.size(), var_name, "/",
+                                 modifiable);
+  }
+}
+
+} // namespace
+
 namespace dolfinx::io::native
 {
 
@@ -43,6 +77,104 @@ dolfinx::mesh::Mesh<T> read_mesh(adios2::IO& io, adios2::Engine& engine,
                                  dolfinx::mesh::GhostMode ghost_mode
                                  = dolfinx::mesh::GhostMode::shared_facet);
 
+/// @brief Write meshtags to a file.
+///
+/// @tparam U float or double
+/// @tparam T ADIOS2 supported type
+/// @param[in] io ADIOS2 IO
+/// @param[in] engine ADIOS2 Engine
+/// @param[in] mesh Mesh of type float or double to write to the file
+/// @param[in] meshtags MeshTags to write to the file
+template <std::floating_point U, typename T>
+void write_meshtags(adios2::IO& io, adios2::Engine& engine,
+                    dolfinx::mesh::Mesh<U>& mesh,
+                    dolfinx::mesh::MeshTags<T>& meshtags)
+{
+  std::string name = meshtags.name;
+  std::uint32_t dim = meshtags.dim();
+
+  spdlog::info("Writing meshtags : {} for entities with dimension: {}", name,
+               dim);
+
+  {
+    adios2::Attribute<std::string> attr_names
+        = define_attr<std::string>(io, "meshtags_names", name);
+    adios2::Attribute<std::uint32_t> attr_dims
+        = define_attr<std::uint32_t>(io, "meshtags_dims", dim);
+  }
+
+  //
+  std::uint32_t num_dofs_per_entity;
+  {
+    const fem::CoordinateElement<U>& cmap = mesh.geometry().cmap();
+    fem::ElementDofLayout geom_layout = cmap.create_dof_layout();
+    num_dofs_per_entity = geom_layout.num_entity_closure_dofs(dim);
+  }
+
+  std::uint64_t num_tag_entities_global, offset;
+
+  // NOTE: For correctness of MPI_ calls, the following should match the
+  // uint64_t type
+  std::uint64_t num_tag_entities_local;
+  std::vector<std::int64_t> array;
+  {
+    std::uint32_t num_entities_local;
+    {
+      std::shared_ptr<const mesh::Topology> topology = mesh.topology();
+      assert(topology);
+
+      num_entities_local = topology->index_map(dim)->size_local();
+    }
+
+    std::span<const std::int32_t> tag_entities = meshtags.indices();
+    assert(std::ranges::is_sorted(tag_entities));
+
+    std::uint32_t num_tag_entities = tag_entities.size();
+
+    num_tag_entities_local
+        = std::upper_bound(tag_entities.begin(), tag_entities.end(),
+                           num_entities_local)
+          - tag_entities.begin();
+
+    // Compute the global offset for owned tagged entities
+    offset = (std::uint64_t)0;
+    MPI_Exscan(&num_tag_entities_local, &offset, 1, MPI_UINT64_T, MPI_SUM,
+               mesh.comm());
+
+    // Compute the global size of tagged entities
+    num_tag_entities_global = (std::uint64_t)0;
+    MPI_Allreduce(&num_tag_entities_local, &num_tag_entities_global, 1,
+                  MPI_UINT64_T, MPI_SUM, mesh.comm());
+
+    std::vector<std::int32_t> entities_to_geometry = mesh::entities_to_geometry(
+        mesh, dim, tag_entities.subspan(0, num_tag_entities_local), false);
+
+    array.resize(entities_to_geometry.size());
+
+    std::iota(array.begin(), array.end(), 0);
+
+    std::shared_ptr<const common::IndexMap> imap = mesh.geometry().index_map();
+    imap->local_to_global(entities_to_geometry, array);
+  }
+
+  std::span<const T> values = meshtags.values();
+  const std::span<const T> local_values(values.begin(), num_tag_entities_local);
+
+  adios2::Variable<std::int64_t> var_topology = io.DefineVariable<std::int64_t>(
+      name + "_topology", {num_tag_entities_global, num_dofs_per_entity},
+      {offset, 0}, {num_tag_entities_local, num_dofs_per_entity},
+      adios2::ConstantDims);
+
+  adios2::Variable<T> var_values = io.DefineVariable<T>(
+      name + "_values", {num_tag_entities_global}, {offset},
+      {num_tag_entities_local}, adios2::ConstantDims);
+
+  engine.BeginStep();
+  engine.Put(var_topology, array.data());
+  engine.Put(var_values, local_values.data());
+  engine.EndStep();
+}
+
 } // namespace dolfinx::io::native
 
 namespace dolfinx::io::impl_native