diff --git a/cpp/demo/checkpointing/main.cpp b/cpp/demo/checkpointing/main.cpp index 06b9ad00d1e..11f309cafe1 100644 --- a/cpp/demo/checkpointing/main.cpp +++ b/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 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 entities_array(num_entities * num_dofs_per_entity); + std::vector entities_offsets(num_entities + 1); + std::uint64_t offset = topo_imap->local_range()[0]; + std::vector 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 entities_local(entities_array, + entities_offsets); + + auto meshtags = std::make_shared>( + mesh::create_meshtags(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(io, engine, *mesh, *meshtags); engine.Close(); } diff --git a/cpp/dolfinx/io/checkpointing.h b/cpp/dolfinx/io/checkpointing.h index b006a926566..52adcb4899a 100644 --- a/cpp/dolfinx/io/checkpointing.h +++ b/cpp/dolfinx/io/checkpointing.h @@ -12,12 +12,46 @@ #include #include #include +#include #include #include /// @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 +adios2::Attribute define_attr(adios2::IO& io, const std::string& name, + T& value, std::string var_name = "") +{ + bool modifiable = true; + if (adios2::Attribute attr = io.InquireAttribute(name); attr) + { + std::vector data = attr.Data(); + data.push_back(value); + return io.DefineAttribute(name, data.data(), data.size(), var_name, "/", + modifiable); + } + else + { + std::vector data{value}; + return io.DefineAttribute(name, data.data(), data.size(), var_name, "/", + modifiable); + } +} + +} // namespace + namespace dolfinx::io::native { @@ -43,6 +77,104 @@ dolfinx::mesh::Mesh 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 +void write_meshtags(adios2::IO& io, adios2::Engine& engine, + dolfinx::mesh::Mesh& mesh, + dolfinx::mesh::MeshTags& meshtags) +{ + std::string name = meshtags.name; + std::uint32_t dim = meshtags.dim(); + + spdlog::info("Writing meshtags : {} for entities with dimension: {}", name, + dim); + + { + adios2::Attribute attr_names + = define_attr(io, "meshtags_names", name); + adios2::Attribute attr_dims + = define_attr(io, "meshtags_dims", dim); + } + + // + std::uint32_t num_dofs_per_entity; + { + const fem::CoordinateElement& 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 array; + { + std::uint32_t num_entities_local; + { + std::shared_ptr topology = mesh.topology(); + assert(topology); + + num_entities_local = topology->index_map(dim)->size_local(); + } + + std::span 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 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 imap = mesh.geometry().index_map(); + imap->local_to_global(entities_to_geometry, array); + } + + std::span values = meshtags.values(); + const std::span local_values(values.begin(), num_tag_entities_local); + + adios2::Variable var_topology = io.DefineVariable( + name + "_topology", {num_tag_entities_global, num_dofs_per_entity}, + {offset, 0}, {num_tag_entities_local, num_dofs_per_entity}, + adios2::ConstantDims); + + adios2::Variable var_values = io.DefineVariable( + 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