diff --git a/cpp/demo/checkpointing/CMakeLists.txt b/cpp/demo/checkpointing/CMakeLists.txt new file mode 100644 index 00000000000..bd43a542bfd --- /dev/null +++ b/cpp/demo/checkpointing/CMakeLists.txt @@ -0,0 +1,40 @@ +# This file was generated by running +# +# python cmake/scripts/generate-cmakefiles.py from dolfinx/cpp +# +cmake_minimum_required(VERSION 3.19) + +set(PROJECT_NAME demo_checkpointing) +project(${PROJECT_NAME} LANGUAGES C CXX) + +# Set C++20 standard +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +if(NOT TARGET dolfinx) + find_package(DOLFINX REQUIRED) +endif() + +set(CMAKE_INCLUDE_CURRENT_DIR ON) + +add_executable(${PROJECT_NAME} main.cpp) +target_link_libraries(${PROJECT_NAME} dolfinx) + +# Do not throw error for 'multi-line comments' (these are typical in rst which +# includes LaTeX) +include(CheckCXXCompilerFlag) +check_cxx_compiler_flag("-Wno-comment" HAVE_NO_MULTLINE) +set_source_files_properties( + main.cpp + PROPERTIES + COMPILE_FLAGS + "$<$:-Wno-comment -Wall -Wextra -pedantic -Werror>" +) + +# Test targets (used by DOLFINx testing system) +set(TEST_PARAMETERS2 -np 2 ${MPIEXEC_PARAMS} "./${PROJECT_NAME}") +set(TEST_PARAMETERS3 -np 3 ${MPIEXEC_PARAMS} "./${PROJECT_NAME}") +add_test(NAME ${PROJECT_NAME}_mpi_2 COMMAND "mpirun" ${TEST_PARAMETERS2}) +add_test(NAME ${PROJECT_NAME}_mpi_3 COMMAND "mpirun" ${TEST_PARAMETERS3}) +add_test(NAME ${PROJECT_NAME}_serial COMMAND ${PROJECT_NAME}) diff --git a/cpp/demo/checkpointing/checkpointing.yml b/cpp/demo/checkpointing/checkpointing.yml new file mode 100644 index 00000000000..3ca79ee71d2 --- /dev/null +++ b/cpp/demo/checkpointing/checkpointing.yml @@ -0,0 +1,10 @@ +--- +# adios2 config.yaml +# IO YAML Sequence (-) Nodes to allow for multiple IO nodes +# IO name referred in code with DeclareIO is mandatory + +- IO: "mesh-write" + + Engine: + # If Type is missing or commented out, default Engine is picked up + Type: "BP5" diff --git a/cpp/demo/checkpointing/main.cpp b/cpp/demo/checkpointing/main.cpp new file mode 100644 index 00000000000..adce3bf94a0 --- /dev/null +++ b/cpp/demo/checkpointing/main.cpp @@ -0,0 +1,146 @@ +// ```text +// Copyright (C) 2024 Abdullah Mujahid, Jørgen S. Dokken, Jack S. Hale +// This file is part of DOLFINx (https://www.fenicsproject.org) +// SPDX-License-Identifier: LGPL-3.0-or-later +// ``` + +// # Checkpointing +// + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace dolfinx; +using namespace dolfinx::io; + +// Create cell meshtags + +template +std::shared_ptr> +create_meshtags(dolfinx::mesh::Mesh& mesh) +{ + // 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::int32_t num_entities = topo_imap->size_local(); + + auto cmap = 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::int32_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)); + + return meshtags; +} + +int main(int argc, char* argv[]) +{ + dolfinx::init_logging(argc, argv); + MPI_Init(&argc, &argv); + + // Create mesh and function space + auto part = mesh::create_cell_partitioner(mesh::GhostMode::shared_facet); + auto mesh = std::make_shared>(mesh::create_rectangle( + MPI_COMM_WORLD, {{{0.0, 0.0}, {1.0, 1.0}}}, {4, 4}, + mesh::CellType::quadrilateral, part)); + + auto meshtags = create_meshtags(*mesh); + + try + { + // Set up ADIOS2 IO and Engine + adios2::ADIOS adios(mesh->comm()); + + adios2::IO io = adios.DeclareIO("mesh-write"); + io.SetEngine("BP5"); + 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(); + } + catch (std::exception& e) + { + std::cout << "ERROR: ADIOS2 exception: " << e.what() << "\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + + try + { + // Set up ADIOS2 IO and Engine + adios2::ADIOS adios_read(MPI_COMM_WORLD); + adios2::IO io_read = adios_read.DeclareIO("mesh-read"); + io_read.SetEngine("BP5"); + adios2::Engine engine_read = io_read.Open("mesh.bp", adios2::Mode::Read); + + engine_read.BeginStep(); + auto mesh_read + = io::native::read_mesh(io_read, engine_read, MPI_COMM_WORLD); + + mesh::MeshTags mt + = io::native::read_meshtags( + io_read, engine_read, mesh_read, "mesh_tags"); + + if (engine_read.BetweenStepPairs()) + { + engine_read.EndStep(); + } + + engine_read.Close(); + + adios2::IO io_write = adios_read.DeclareIO("mesh-write"); + io_write.SetEngine("BP5"); + adios2::Engine engine_write + = io_write.Open("mesh2.bp", adios2::Mode::Write); + + io::native::write_mesh(io_write, engine_write, mesh_read); + io::native::write_meshtags(io_write, engine_write, + mesh_read, mt); + engine_write.Close(); + } + catch (std::exception& e) + { + std::cout << "ERROR: ADIOS2 exception: " << e.what() << "\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + + MPI_Finalize(); + return 0; +} diff --git a/cpp/doc/source/demo.rst b/cpp/doc/source/demo.rst index 7ce9d8d80db..96b21104885 100644 --- a/cpp/doc/source/demo.rst +++ b/cpp/doc/source/demo.rst @@ -42,3 +42,4 @@ Experimental :maxdepth: 1 demos/demo_mixed_topology.md + demos/demo_checkpointing.md diff --git a/cpp/dolfinx/io/ADIOS2_utils.h b/cpp/dolfinx/io/ADIOS2_utils.h new file mode 100644 index 00000000000..297dd38ab73 --- /dev/null +++ b/cpp/dolfinx/io/ADIOS2_utils.h @@ -0,0 +1,104 @@ +// Copyright (C) 2024 Abdullah Mujahid, Jørgen S. Dokken and Garth N. Wells +// +// This file is part of DOLFINX (https://www.fenicsproject.org) +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#pragma once + +#ifdef HAS_ADIOS2 + +#include +#include +#include +#include + +/// @file ADIOS2_utils.h +/// @brief Utils for ADIOS2 + +namespace +{ +std::map string_to_mode{ + {"write", adios2::Mode::Write}, + {"read", adios2::Mode::Read}, + {"append", adios2::Mode::Append}, +}; +} // namespace + +namespace dolfinx::io +{ + +/// ADIOS2-based writers/readers +class ADIOS2Wrapper +{ +public: + /// @brief Create an ADIOS2-based engine writer/reader + /// @param[in] comm The MPI communicator + /// @param[in] filename Name of output file + /// @param[in] tag The ADIOS2 IO name + /// @param[in] engine_type ADIOS2 engine type. See + /// https://adios2.readthedocs.io/en/latest/engines/engines.html. + /// @param[in] mode ADIOS2 mode, default is Write or Read + ADIOS2Wrapper(MPI_Comm comm, std::string filename, std::string tag, + std::string engine_type = "BP5", std::string mode = "write") + : _adios(std::make_shared(comm)), + _io(std::make_shared(_adios->DeclareIO(tag))) + { + _io->SetEngine(engine_type); + _engine = std::make_shared( + _io->Open(filename, string_to_mode[mode])); + } + + /// @brief Create an ADIOS2-based engine writer/reader + /// @param[in] config_file Path to config file to set up ADIOS2 engine from + /// @param[in] comm The MPI communicator + /// @param[in] filename Name of output file + /// @param[in] tag The ADIOS2 IO name + /// @param[in] mode ADIOS2 mode, default is Write or Read + ADIOS2Wrapper(std::string config_file, MPI_Comm comm, std::string filename, + std::string tag, std::string mode = "append") + { + _adios = std::make_shared(config_file, comm); + _io = std::make_shared(_adios->DeclareIO(tag)); + _engine = std::make_shared( + _io->Open(filename, string_to_mode[mode])); + } + + /// @brief Move constructor + ADIOS2Wrapper(ADIOS2Wrapper&& ADIOS2) = default; + + /// @brief Copy constructor + ADIOS2Wrapper(const ADIOS2Wrapper&) = delete; + + /// @brief Destructor + ~ADIOS2Wrapper() { close(); } + + /// @brief Move assignment + ADIOS2Wrapper& operator=(ADIOS2Wrapper&& ADIOS2) = default; + + // Copy assignment + ADIOS2Wrapper& operator=(const ADIOS2Wrapper&) = delete; + + /// @brief Close the file + void close() + { + assert(_engine); + if (*_engine) + _engine->Close(); + } + + /// @brief Get the IO object + std::shared_ptr io() { return _io; } + + /// @brief Get the Engine object + std::shared_ptr engine() { return _engine; } + +private: + std::shared_ptr _adios; + std::shared_ptr _io; + std::shared_ptr _engine; +}; + +} // namespace dolfinx::io + +#endif diff --git a/cpp/dolfinx/io/CMakeLists.txt b/cpp/dolfinx/io/CMakeLists.txt index a3060dd6852..c03d18de19f 100644 --- a/cpp/dolfinx/io/CMakeLists.txt +++ b/cpp/dolfinx/io/CMakeLists.txt @@ -1,7 +1,9 @@ set(HEADERS_io ${CMAKE_CURRENT_SOURCE_DIR}/dolfinx_io.h + ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2_utils.h ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.h ${CMAKE_CURRENT_SOURCE_DIR}/cells.h + ${CMAKE_CURRENT_SOURCE_DIR}/checkpointing.h ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.h ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.h ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.h @@ -16,6 +18,7 @@ target_sources( dolfinx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cells.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/checkpointing.cpp ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.cpp ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.cpp ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.cpp diff --git a/cpp/dolfinx/io/checkpointing.cpp b/cpp/dolfinx/io/checkpointing.cpp new file mode 100644 index 00000000000..8fbed23a51e --- /dev/null +++ b/cpp/dolfinx/io/checkpointing.cpp @@ -0,0 +1,439 @@ +// Copyright (C) 2024 Abdullah Mujahid, Jørgen S. Dokken, Jack S. Hale +// +// This file is part of DOLFINX (https://www.fenicsproject.org) +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#ifdef HAS_ADIOS2 + +#include "checkpointing.h" +#include +#include +#include +#include +#include +#include +#include + +/// @file checkpointing.h +/// @brief ADIOS2 based checkpointing +namespace +{ +/// basix enum to string +std::map variant_to_string{ + {basix::element::lagrange_variant::unset, "unset"}, + {basix::element::lagrange_variant::equispaced, "equispaced"}, + {basix::element::lagrange_variant::gll_warped, "gll_warped"}, + {basix::element::lagrange_variant::gll_isaac, "gll_isaac"}, + {basix::element::lagrange_variant::gll_centroid, "gll_centroid"}, + {basix::element::lagrange_variant::chebyshev_warped, "chebyshev_warped"}, + {basix::element::lagrange_variant::chebyshev_isaac, "chebyshev_isaac"}, + {basix::element::lagrange_variant::gl_warped, "gl_warped"}, + {basix::element::lagrange_variant::gl_isaac, "gl_isaac"}, + {basix::element::lagrange_variant::gl_centroid, "gl_centroid"}, + {basix::element::lagrange_variant::legendre, "legendre"}, + {basix::element::lagrange_variant::bernstein, "bernstein"}, +}; + +/// string to basix enum +std::map string_to_variant{ + {"unset", basix::element::lagrange_variant::unset}, + {"equispaced", basix::element::lagrange_variant::equispaced}, + {"gll_warped", basix::element::lagrange_variant::gll_warped}, + {"gll_isaac", basix::element::lagrange_variant::gll_isaac}, + {"gll_centroid", basix::element::lagrange_variant::gll_centroid}, + {"chebyshev_warped", basix::element::lagrange_variant::chebyshev_warped}, + {"chebyshev_isaac", basix::element::lagrange_variant::chebyshev_isaac}, + {"gl_warped", basix::element::lagrange_variant::gl_warped}, + {"gl_isaac", basix::element::lagrange_variant::gl_isaac}, + {"gl_centroid", basix::element::lagrange_variant::gl_centroid}, + {"legendre", basix::element::lagrange_variant::legendre}, + {"bernstein", basix::element::lagrange_variant::bernstein}, +}; + +} // namespace + +namespace dolfinx::io::native +{ +//----------------------------------------------------------------------------- +template +void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh) +{ + + const mesh::Geometry& geometry = mesh.geometry(); + std::shared_ptr topology = mesh.topology(); + assert(topology); + + // Variables/attributes to save + std::int32_t dim = geometry.dim(); + std::int32_t tdim = topology->dim(); + std::uint64_t num_nodes_global, offset, num_cells_global, cell_offset; + std::uint32_t num_nodes_local, num_cells_local, degree; + std::string cell_type, lagrange_variant; + std::vector array_global; + std::vector offsets_global; + + std::shared_ptr geom_imap; + + // Nodes information + { + geom_imap = geometry.index_map(); + num_nodes_global = geom_imap->size_global(); + num_nodes_local = geom_imap->size_local(); + offset = geom_imap->local_range()[0]; + } + + // Cells information + { + const std::shared_ptr topo_imap + = topology->index_map(tdim); + assert(topo_imap); + + num_cells_global = topo_imap->size_global(); + num_cells_local = topo_imap->size_local(); + cell_offset = topo_imap->local_range()[0]; + } + + // Coordinate element information + { + const fem::CoordinateElement& cmap = geometry.cmap(); + cell_type = mesh::to_string(cmap.cell_shape()); + degree = cmap.degree(); + lagrange_variant = variant_to_string[cmap.variant()]; + } + + const std::span mesh_x = geometry.x(); + + std::uint32_t num_dofs_per_cell; + // Connectivity + { + auto dofmap = geometry.dofmap(); + num_dofs_per_cell = dofmap.extent(1); + std::vector connectivity; + connectivity.reserve(num_cells_local * num_dofs_per_cell); + for (std::size_t i = 0; i < num_cells_local; ++i) + for (std::size_t j = 0; j < num_dofs_per_cell; ++j) + connectivity.push_back(dofmap(i, j)); + + array_global.resize(num_cells_local * num_dofs_per_cell); + std::iota(array_global.begin(), array_global.end(), 0); + + geom_imap->local_to_global(connectivity, array_global); + offsets_global.resize(num_cells_local + 1); + + // FIXME: use std::ranges::transform + for (std::size_t i = 0; i < num_cells_local + 1; ++i) + offsets_global[i] = (i + cell_offset) * num_dofs_per_cell; + } + + // ADIOS2 write attributes and variables + { + io.DefineAttribute("version", dolfinx::version()); + io.DefineAttribute("git_hash", dolfinx::git_commit_hash()); + io.DefineAttribute("name", mesh.name); + io.DefineAttribute("dim", dim); + io.DefineAttribute("tdim", tdim); + io.DefineAttribute("cell_type", cell_type); + io.DefineAttribute("degree", degree); + io.DefineAttribute("lagrange_variant", lagrange_variant); + + adios2::Variable var_num_nodes + = io.DefineVariable("num_nodes"); + adios2::Variable var_num_cells + = io.DefineVariable("num_cells"); + + adios2::Variable var_x + = io.DefineVariable("x", {num_nodes_global, 3}, {offset, 0}, + {num_nodes_local, 3}, adios2::ConstantDims); + + adios2::Variable var_topology_array + = io.DefineVariable( + "topology_array", {num_cells_global * num_dofs_per_cell}, + {cell_offset * num_dofs_per_cell}, + {num_cells_local * num_dofs_per_cell}, adios2::ConstantDims); + + adios2::Variable var_topology_offsets + = io.DefineVariable( + "topology_offsets", {num_cells_global + 1}, {cell_offset}, + {num_cells_local + 1}, adios2::ConstantDims); + + assert(!engine.BetweenStepPairs()); + engine.BeginStep(); + engine.Put(var_num_nodes, num_nodes_global); + engine.Put(var_num_cells, num_cells_global); + engine.Put(var_x, mesh_x.subspan(0, num_nodes_local * 3).data()); + engine.Put(var_topology_array, array_global.data()); + engine.Put(var_topology_offsets, offsets_global.data()); + engine.EndStep(); + + spdlog::info("Mesh written"); + } +} + +//----------------------------------------------------------------------------- +/// @cond +template void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh); + +template void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh); + +/// @endcond + +//----------------------------------------------------------------------------- +template +dolfinx::mesh::Mesh read_mesh(adios2::IO& io, adios2::Engine& engine, + MPI_Comm comm, + dolfinx::mesh::GhostMode ghost_mode) +{ + + int rank, size; + MPI_Comm_rank(comm, &rank); + MPI_Comm_size(comm, &size); + + if (!engine.BetweenStepPairs()) + { + engine.BeginStep(); + } + + // Compatibility check for version and git commit hash + { + adios2::Attribute var_version + = io.InquireAttribute("version"); + adios2::Attribute var_hash + = io.InquireAttribute("git_hash"); + std::string version = var_version.Data()[0]; + if (version != dolfinx::version()) + { + throw std::runtime_error("Reading from version: " + dolfinx::version() + + " written with version: " + version); + } + + std::string git_hash = var_hash.Data()[0]; + if (git_hash != dolfinx::git_commit_hash()) + { + throw std::runtime_error("Reading from GIT_COMMIT_HASH: " + + dolfinx::git_commit_hash() + + " written with GIT_COMMIT_HASH: " + git_hash); + } + } + // Attributes + std::string name; + std::int32_t dim, tdim, degree; + mesh::CellType cell_type; + basix::element::lagrange_variant lagrange_variant; + + // Read attributes + { + adios2::Attribute var_name + = io.InquireAttribute("name"); + adios2::Attribute var_dim + = io.InquireAttribute("dim"); + adios2::Attribute var_tdim + = io.InquireAttribute("tdim"); + adios2::Attribute var_cell_type + = io.InquireAttribute("cell_type"); + adios2::Attribute var_degree + = io.InquireAttribute("degree"); + adios2::Attribute var_variant + = io.InquireAttribute("lagrange_variant"); + + name = var_name.Data()[0]; + dim = var_dim.Data()[0]; + tdim = var_tdim.Data()[0]; + cell_type = mesh::to_type(var_cell_type.Data()[0]); + degree = var_degree.Data()[0]; + lagrange_variant = string_to_variant[var_variant.Data()[0]]; + } + + spdlog::info( + "Reading mesh with geometric dimension: {} and topological dimension: {}", + dim, tdim); + + // Scalar variables + std::uint64_t num_nodes_global, num_cells_global; + + // Read scalar variables + { + adios2::Variable var_num_nodes + = io.InquireVariable("num_nodes"); + adios2::Variable var_num_cells + = io.InquireVariable("num_cells"); + + engine.Get(var_num_nodes, num_nodes_global); + engine.Get(var_num_cells, num_cells_global); + } + + auto [x_reduced, x_shape] = dolfinx::io::impl_native::read_geometry_data( + io, engine, dim, num_nodes_global, rank, size); + + std::vector array = dolfinx::io::impl_native::read_topology_data( + io, engine, num_cells_global, rank, size); + + assert(engine.BetweenStepPairs()); + engine.EndStep(); + + fem::CoordinateElement element + = fem::CoordinateElement(cell_type, degree, lagrange_variant); + + auto part = mesh::create_cell_partitioner(ghost_mode); + + if (size == 1) + part = nullptr; + + mesh::Mesh mesh = mesh::create_mesh(comm, comm, array, element, comm, + x_reduced, x_shape, part); + + mesh.name = name; + + return mesh; +} + +//----------------------------------------------------------------------------- +/// @cond +template dolfinx::mesh::Mesh +read_mesh(adios2::IO& io, adios2::Engine& engine, MPI_Comm comm, + dolfinx::mesh::GhostMode ghost_mode); + +template dolfinx::mesh::Mesh +read_mesh(adios2::IO& io, adios2::Engine& engine, MPI_Comm comm, + dolfinx::mesh::GhostMode ghost_mode); + +/// @endcond + +} // namespace dolfinx::io::native + +namespace dolfinx::io::impl_native +{ +//----------------------------------------------------------------------------- +std::pair get_counters(int rank, std::uint64_t N, + int size) +{ + assert(rank >= 0); + assert(size > 0); + + // Compute number of items per rank and remainder + const std::uint64_t n = N / size; + const std::uint64_t r = N % size; + + // Compute local range + if (static_cast(rank) < r) + return {rank * (n + 1), n + 1}; + else + return {rank * n + r, n}; +} + +//----------------------------------------------------------------------------- +template +std::pair, std::array> +read_geometry_data(adios2::IO& io, adios2::Engine& engine, int dim, + std::uint64_t num_nodes_global, int rank, int size) +{ + auto [nodes_offset, num_nodes_local] + = dolfinx::io::impl_native::get_counters(rank, num_nodes_global, size); + std::vector x(num_nodes_local * 3); + adios2::Variable var_x = io.InquireVariable("x"); + if (var_x) + { + var_x.SetSelection({{nodes_offset, 0}, {num_nodes_local, 3}}); + engine.Get(var_x, x.data(), adios2::Mode::Sync); + } + else + { + throw std::runtime_error("Coordinates data not found"); + } + + // FIXME: Use std::ranges::transform + std::vector x_reduced(num_nodes_local * dim); + for (std::uint32_t i = 0; i < num_nodes_local; ++i) + { + for (std::uint32_t j = 0; j < (std::uint32_t)dim; ++j) + x_reduced[i * dim + j] = x[i * 3 + j]; + } + + std::array xshape = {num_nodes_local, (std::uint32_t)dim}; + + return {std::move(x_reduced), xshape}; +} + +//----------------------------------------------------------------------------- +std::vector read_topology_data(adios2::IO& io, adios2::Engine& engine, + std::uint64_t num_cells_global, + int rank, int size) +{ + auto [cells_offset, num_cells_local] + = dolfinx::io::impl_native::get_counters(rank, num_cells_global, size); + std::vector offsets(num_cells_local + 1); + + adios2::Variable var_topology_array + = io.InquireVariable("topology_array"); + + adios2::Variable var_topology_offsets + = io.InquireVariable("topology_offsets"); + + if (var_topology_offsets) + { + var_topology_offsets.SetSelection({{cells_offset}, {num_cells_local + 1}}); + engine.Get(var_topology_offsets, offsets.data(), adios2::Mode::Sync); + } + else + { + throw std::runtime_error("Topology offsets not found"); + } + + std::uint64_t count + = static_cast(offsets[num_cells_local] - offsets[0]); + std::vector array(count); + + if (var_topology_array) + { + var_topology_array.SetSelection( + {{static_cast(offsets[0])}, {count}}); + engine.Get(var_topology_array, array.data(), adios2::Mode::Sync); + } + else + { + throw std::runtime_error("Topology array not found"); + } + + return array; +} + +//----------------------------------------------------------------------------- +std::variant, dolfinx::mesh::Mesh> +read_mesh_variant(adios2::IO& io, adios2::Engine& engine, MPI_Comm comm, + dolfinx::mesh::GhostMode ghost_mode) +{ + engine.BeginStep(); + std::string floating_point = io.VariableType("x"); + + if (floating_point == "float") + { + dolfinx::mesh::Mesh mesh + = dolfinx::io::native::read_mesh(io, engine, comm, ghost_mode); + if (engine.BetweenStepPairs()) + { + engine.EndStep(); + } + return mesh; + } + else if (floating_point == "double") + { + dolfinx::mesh::Mesh mesh + = dolfinx::io::native::read_mesh(io, engine, comm, ghost_mode); + if (engine.BetweenStepPairs()) + { + engine.EndStep(); + } + return mesh; + } + else + { + throw std::runtime_error("Floating point type is neither float nor double"); + } +} + +} // namespace dolfinx::io::impl_native + +#endif diff --git a/cpp/dolfinx/io/checkpointing.h b/cpp/dolfinx/io/checkpointing.h new file mode 100644 index 00000000000..c8451123a0e --- /dev/null +++ b/cpp/dolfinx/io/checkpointing.h @@ -0,0 +1,386 @@ +// Copyright (C) 2024 Abdullah Mujahid, Jørgen S. Dokken, Jack S. Hale +// +// This file is part of DOLFINX (https://www.fenicsproject.org) +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#pragma once + +#ifdef HAS_ADIOS2 + +#include "ADIOS2_utils.h" +#include "xdmf_utils.h" +#include +#include +#include +#include +#include +#include +#include +#include + +/// @file checkpointing.h +/// @brief ADIOS2 based checkpointing + +namespace +{ + +template +using mdspan_t = MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan< + const T, MDSPAN_IMPL_STANDARD_NAMESPACE::dextents>; + +/// @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::impl_native +{ +/// @brief Find offset and size. +/// +/// @param[in] rank MPI rank +/// @param[in] N size of data to distribute +/// @param[in] size MPI size +/// @return start and count +std::pair get_counters(int rank, std::uint64_t N, + int size); + +/// @brief Read geometry data +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @param[in] dim The geometric dimension (`0 < dim <= 3`). +/// @param[in] num_nodes_global size of the global array of nodes +/// @param[in] rank MPI rank +/// @param[in] size MPI size +/// @return The point coordinates of row-major storage and +/// itsshape `(num_nodes_local, dim)` +template +std::pair, std::array> +read_geometry_data(adios2::IO& io, adios2::Engine& engine, int dim, + std::uint64_t num_nodes_global, int rank, int size); + +/// @brief Read topology array +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @param[in] num_cells_global global number of cells +/// @param[in] rank MPI rank +/// @param[in] size MPI size +/// @return The cell-to-node connectivity in a flattened array +std::vector read_topology_data(adios2::IO& io, adios2::Engine& engine, + std::uint64_t num_cells_global, + int rank, int size); + +/// @brief Read mesh from a file. +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @param[in] comm comm +/// @param[in] ghost_mode The requested type of cell ghosting/overlap +/// @return mesh reconstructed from the data +std::variant, dolfinx::mesh::Mesh> +read_mesh_variant(adios2::IO& io, adios2::Engine& engine, + MPI_Comm comm = MPI_COMM_WORLD, + dolfinx::mesh::GhostMode ghost_mode + = dolfinx::mesh::GhostMode::shared_facet); + +} // namespace dolfinx::io::impl_native + +namespace dolfinx::io::native +{ + +/// @brief Write mesh to a file. +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @param[in] mesh Mesh of type float or double to write to the file +template +void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh); + +/// @brief Read mesh from a file. +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @param[in] comm comm +/// @param[in] ghost_mode The requested type of cell ghosting/overlap +/// @return mesh reconstructed from the data +template +dolfinx::mesh::Mesh read_mesh(adios2::IO& io, adios2::Engine& engine, + MPI_Comm comm = MPI_COMM_WORLD, + 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(); + std::string dtype; + + if (std::is_same_v) + { + dtype = "int32_t"; + } + else if (std::is_same_v) + { + dtype = "double"; + } + else + { + throw std::runtime_error("The datatype associated with the meshtags values " + "is not supported yet"); + } + + spdlog::info( + "Writing meshtags : {} for entities with dimension: {} of data type: {}", + name, dim, dtype); + + { + adios2::Attribute attr_names + = define_attr(io, "meshtags_names", name); + adios2::Attribute attr_dtypes + = define_attr(io, "meshtags_dtypes", dtype); + 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; + { + // Fetch number of local (owned) entities + std::uint32_t num_entities_local; + { + std::shared_ptr topology = mesh.topology(); + assert(topology); + + topology->create_entities(dim); + 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(); +} + +/// @brief Read meshtags from 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 which meshtags are +/// associated +/// @param[in] name name of the meshtags +/// @return meshtags +template +dolfinx::mesh::MeshTags read_meshtags(adios2::IO& io, adios2::Engine& engine, + dolfinx::mesh::Mesh& mesh, + std::string name) +{ + + int rank, size; + MPI_Comm_rank(mesh.comm(), &rank); + MPI_Comm_size(mesh.comm(), &size); + + if (!engine.BetweenStepPairs()) + { + engine.BeginStep(); + } + + // Attributes + std::int32_t dim; + // Check if meshtags of given name exists, and find the dim + { + adios2::Attribute var_names + = io.InquireAttribute("meshtags_names"); + adios2::Attribute var_dims + = io.InquireAttribute("meshtags_dims"); + + std::vector names = var_names.Data(); + std::vector dims = var_dims.Data(); + auto pos = std::ranges::find(names, name); + if (pos != names.end()) + { + auto it = std::ranges::distance(names.begin(), pos); + dim = dims[it]; + } + else + { + throw std::runtime_error("Meshtags : " + name + " not found"); + } + } + + spdlog::info("Reading meshtags : {} for entities with dimension: {}", name, + dim); + + // Read entities topology + adios2::Variable var_topology_array + = io.InquireVariable(name + "_topology"); + std::vector shape = var_topology_array.Shape(); + + auto [offset, num_tag_entities_local] + = dolfinx::io::impl_native::get_counters(rank, shape[0], size); + + std::uint64_t num_dofs_per_entity = shape[1]; + std::vector array(num_tag_entities_local * num_dofs_per_entity); + var_topology_array.SetSelection( + {{offset, 0}, {num_tag_entities_local, num_dofs_per_entity}}); + engine.Get(var_topology_array, array.data(), adios2::Mode::Sync); + + // Read entities tagged values + adios2::Variable var_values = io.InquireVariable(name + "_values"); + std::vector values(num_tag_entities_local); + var_values.SetSelection({{offset}, {num_tag_entities_local}}); + engine.Get(var_values, values.data(), adios2::Mode::Sync); + + // Redistribute data + std::pair, std::vector> entities_values; + std::int64_t num_entities, num_vert_per_entity; + std::shared_ptr topology = mesh.topology(); + assert(topology); + + { + mdspan_t entities_span( + array.data(), num_tag_entities_local, num_dofs_per_entity); + + const mesh::Geometry& geometry = mesh.geometry(); + + auto xdofmap_span = geometry.dofmap(); + const std::vector& input_global_indices + = geometry.input_global_indices(); + + const fem::CoordinateElement& cmap = geometry.cmap(); + const fem::ElementDofLayout cmap_dof_layout = cmap.create_dof_layout(); + std::int64_t num_nodes_g = geometry.index_map()->size_global(); + + std::span input_global_indices_span( + input_global_indices.data(), input_global_indices.size()); + std::span values_span(values.data(), values.size()); + + entities_values = dolfinx::io::xdmf_utils::distribute_entity_data( + *topology, input_global_indices_span, num_nodes_g, cmap_dof_layout, + xdofmap_span, dim, entities_span, values_span); + + num_vert_per_entity = dolfinx::mesh::cell_num_entities( + dolfinx::mesh::cell_entity_type(topology->cell_type(), dim, 0), 0); + + num_entities = entities_values.first.size() / num_vert_per_entity; + } + + // Construct MeshTags + std::vector offsets(num_entities + 1); + std::ranges::transform(offsets, offsets.begin(), [&, i = 0](auto e) mutable + { return (i++) * num_vert_per_entity; }); + + dolfinx::graph::AdjacencyList entities(entities_values.first, offsets); + std::span values_span(entities_values.second.data(), + entities_values.second.size()); + + topology->create_connectivity(dim, 0); + topology->create_connectivity(dim, topology->dim()); + + dolfinx::mesh::MeshTags meshtags + = dolfinx::mesh::create_meshtags(topology, dim, entities, values_span); + + meshtags.name = name; + + return meshtags; +} + +} // namespace dolfinx::io::native + +#endif diff --git a/cpp/dolfinx/io/dolfinx_io.h b/cpp/dolfinx/io/dolfinx_io.h index 551fecffdcf..1bb24940950 100644 --- a/cpp/dolfinx/io/dolfinx_io.h +++ b/cpp/dolfinx/io/dolfinx_io.h @@ -9,5 +9,7 @@ namespace dolfinx::io // DOLFINx io interface +#include #include #include +#include diff --git a/python/demo/checkpointing.yml b/python/demo/checkpointing.yml new file mode 100644 index 00000000000..3ca79ee71d2 --- /dev/null +++ b/python/demo/checkpointing.yml @@ -0,0 +1,10 @@ +--- +# adios2 config.yaml +# IO YAML Sequence (-) Nodes to allow for multiple IO nodes +# IO name referred in code with DeclareIO is mandatory + +- IO: "mesh-write" + + Engine: + # If Type is missing or commented out, default Engine is picked up + Type: "BP5" diff --git a/python/demo/demo_checkpointing.py b/python/demo/demo_checkpointing.py new file mode 100644 index 00000000000..255fd27d701 --- /dev/null +++ b/python/demo/demo_checkpointing.py @@ -0,0 +1,105 @@ +# --- +# jupyter: +# jupytext: +# text_representation: +# extension: .py +# format_name: light +# format_version: '1.5' +# jupytext_version: 1.13.6 +# --- + +# # Checkpointing +# +# This demo is implemented in {download}`demo_checkpointing.py`. It +# illustrates checkpointing using ADIOS2: +# + +# + +import os + +import dolfinx + +if not dolfinx.common.has_adios2: + print("This demo requires DOLFINx to be compiled with ADIOS2 enabled.") + exit(0) + +from mpi4py import MPI + +import numpy as np + +from dolfinx import io, mesh + +# - + +# Note that it is important to first `from mpi4py import MPI` to +# ensure that MPI is correctly initialised. + +# We create a rectangular {py:class}`Mesh ` using +# {py:func}`create_rectangle `, and +# save it to a file `mesh.bp` + +# + +msh = mesh.create_rectangle( + comm=MPI.COMM_WORLD, + points=((0.0, 0.0), (1.0, 1.0)), + n=(8, 8), + cell_type=mesh.CellType.triangle, +) +# - + +# Create meshtags of all entity types for the above mesh + +# + +tags = {} +for dim in range(msh.topology.dim + 1): + msh.topology.create_connectivity(dim, msh.topology.dim) + imap = msh.topology.index_map(dim) + num_entities_local = imap.size_local + entities = np.arange(num_entities_local, dtype=np.int32) + values = imap.local_range[0] + entities + mt = mesh.meshtags(msh, dim, entities, values) + mt.name = f"entity_{dim}" + tags[mt.name] = mt + +# - + +# + +config_path = os.getcwd() + "/checkpointing.yml" +adios2 = io.ADIOS2(config_path, msh.comm, filename="mesh.bp", tag="mesh-write", mode="write") +# - + +# + +io.write_mesh(adios2, msh) +for mt_name, mt in tags.items(): + io.write_meshtags(adios2, msh, mt) + +adios2.close() +# - + +# + +adios2_read = io.ADIOS2( + msh.comm, filename="mesh.bp", tag="mesh-read", engine_type="BP5", mode="read" +) +# - + +# + +msh_read = io.read_mesh(adios2_read, msh.comm) +tags_read = io.read_meshtags(adios2_read, msh_read) +adios2_read.close() +# - + +# + +adios2_write = io.ADIOS2( + msh_read.comm, filename="mesh2.bp", tag="mesh-write", engine_type="BP5", mode="write" +) +# - + +# + +io.write_mesh(adios2_write, msh_read) + +for mt_name, mt_read in tags_read.items(): + io.write_meshtags(adios2_write, msh_read, mt_read) + + +adios2_write.close() +# - diff --git a/python/doc/source/demos.rst b/python/doc/source/demos.rst index bbbcb889137..5326d08a6c1 100644 --- a/python/doc/source/demos.rst +++ b/python/doc/source/demos.rst @@ -65,6 +65,7 @@ Interpolation, IO and visualisation demos/demo_pyvista.md demos/demo_interpolation-io.md + demos/demo_checkpointing.md Advanced iterative solvers @@ -103,6 +104,7 @@ List of all demos demos/demo_static-condensation.md demos/demo_pyvista.md demos/demo_interpolation-io.md + demos/demo_checkpointing.md demos/demo_types.md demos/demo_lagrange_variants.md demos/demo_tnt-elements.md diff --git a/python/dolfinx/__init__.py b/python/dolfinx/__init__.py index dcdfa35aa3f..b442d99258d 100644 --- a/python/dolfinx/__init__.py +++ b/python/dolfinx/__init__.py @@ -30,10 +30,11 @@ from dolfinx.common import ( TimingType, git_commit_hash, + has_adios2, has_debug, has_kahip, - has_petsc, has_parmetis, + has_petsc, list_timings, timing, ) @@ -70,9 +71,11 @@ def get_include(user=False): "utils", "TimingType", "git_commit_hash", + "has_adios2", "has_debug", "has_kahip", "has_parmetis", + "has_petsc", "list_timings", "timing", ] diff --git a/python/dolfinx/io/__init__.py b/python/dolfinx/io/__init__.py index 3a53be8d7eb..2dd142d731e 100644 --- a/python/dolfinx/io/__init__.py +++ b/python/dolfinx/io/__init__.py @@ -13,6 +13,27 @@ if _cpp.common.has_adios2: # FidesWriter and VTXWriter require ADIOS2 - from dolfinx.io.utils import FidesMeshPolicy, FidesWriter, VTXMeshPolicy, VTXWriter + from dolfinx.io.utils import ( + ADIOS2, + FidesMeshPolicy, + FidesWriter, + VTXMeshPolicy, + VTXWriter, + read_mesh, + read_meshtags, + write_mesh, + write_meshtags, + ) - __all__ = [*__all__, "FidesWriter", "VTXWriter", "FidesMeshPolicy", "VTXMeshPolicy"] + __all__ = [ + *__all__, + "FidesWriter", + "VTXWriter", + "FidesMeshPolicy", + "VTXMeshPolicy", + "ADIOS2", + "write_mesh", + "write_meshtags", + "read_mesh", + "read_meshtags", + ] diff --git a/python/dolfinx/io/utils.py b/python/dolfinx/io/utils.py index a479f898998..4b300c10b50 100644 --- a/python/dolfinx/io/utils.py +++ b/python/dolfinx/io/utils.py @@ -23,7 +23,13 @@ from dolfinx.fem import Function from dolfinx.mesh import GhostMode, Mesh, MeshTags -__all__ = ["VTKFile", "XDMFFile", "cell_perm_gmsh", "cell_perm_vtk", "distribute_entity_data"] +__all__ = [ + "VTKFile", + "XDMFFile", + "cell_perm_gmsh", + "cell_perm_vtk", + "distribute_entity_data", +] def _extract_cpp_objects(functions: typing.Union[Mesh, Function, tuple[Function], list[Function]]): @@ -38,7 +44,18 @@ def _extract_cpp_objects(functions: typing.Union[Mesh, Function, tuple[Function] if _cpp.common.has_adios2: from dolfinx.cpp.io import FidesMeshPolicy, VTXMeshPolicy # F401 - __all__ = [*__all__, "FidesWriter", "VTXWriter", "FidesMeshPolicy", "VTXMeshPolicy"] + __all__ = [ + *__all__, + "FidesWriter", + "VTXWriter", + "FidesMeshPolicy", + "VTXMeshPolicy", + "ADIOS2", + "read_mesh", + "read_meshtags", + "write_mesh", + "write_meshtags", + ] class VTXWriter: """Writer for VTX files, using ADIOS2 to create the files. @@ -186,6 +203,72 @@ def write(self, t: float): def close(self): self._cpp_object.close() + class ADIOS2(_cpp.io.ADIOS2): + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, traceback): + self.close() + + def write_mesh(ADIOS2: ADIOS2, mesh: Mesh) -> None: + """Write mesh to a file using ADIOS2""" + dtype = mesh.geometry.x.dtype # type: ignore + if np.issubdtype(dtype, np.float32): + _writer = _cpp.io.write_mesh_float32 + elif np.issubdtype(dtype, np.float64): + _writer = _cpp.io.write_mesh_float64 + + return _writer(ADIOS2, mesh._cpp_object) + + def read_mesh( + ADIOS2: ADIOS2, comm: _MPI.Comm, ghost_mode: GhostMode = GhostMode.shared_facet + ) -> Mesh: + """Read mesh from a file using ADIOS2""" + msh = _cpp.io.read_mesh(ADIOS2, comm, ghost_mode) + + element = basix.ufl.element( + basix.ElementFamily.P, + basix.CellType[msh.topology.cell_name()], + msh.geometry.cmap.degree, + basix.LagrangeVariant(int(msh.geometry.cmap.variant)), + shape=(msh.geometry.x.shape[1],), + dtype=msh.geometry.x.dtype, + ) + domain = ufl.Mesh(element) + + return Mesh(msh, domain) + + def write_meshtags(ADIOS2: ADIOS2, mesh: Mesh, meshtags: MeshTags) -> None: + """ + Write meshtags to a file using ADIOS2 + + Args: + ADIOS2: Wrapper to ADIOS2 + mesh: The mesh + meshtags: The associated meshtags + """ + + return _cpp.io.write_meshtags(ADIOS2, mesh._cpp_object, meshtags._cpp_object) + + def read_meshtags(ADIOS2: ADIOS2, mesh: Mesh) -> dict[str, MeshTags]: + """ + Read all meshtags from a file using ADIOS2 + + Args: + ADIOS2: Wrapper to ADIOS2 + mesh: The mesh + Returns: + meshtags: All meshtags stored in the file are returned as a dictionary + with (meshtags_name, meshtags) key-value pairs. + """ + + tags = _cpp.io.read_meshtags(ADIOS2, mesh._cpp_object) + + for name, mt_cpp in tags.items(): + tags[name] = MeshTags(mt_cpp) + + return tags + class VTKFile(_cpp.io.VTKFile): """Interface to VTK files. diff --git a/python/dolfinx/wrappers/io.cpp b/python/dolfinx/wrappers/io.cpp index ed7ae36816b..8002bb0365d 100644 --- a/python/dolfinx/wrappers/io.cpp +++ b/python/dolfinx/wrappers/io.cpp @@ -11,9 +11,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -24,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -232,10 +235,151 @@ void declare_data_types(nb::module_& m) nb::arg("values").noconvert()); } +#ifdef HAS_ADIOS2 +template +void declare_write_mesh(nb::module_& m, std::string type) +{ + // dolfinx::io::native::write_mesh + std::string pyfunction_write_mesh_name = std::string("write_mesh_") + type; + m.def( + pyfunction_write_mesh_name.c_str(), + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, dolfinx::mesh::Mesh& mesh) + { + auto io = ADIOS2.io(); + auto engine = ADIOS2.engine(); + return dolfinx::io::native::write_mesh(*io, *engine, mesh); + }, + nb::arg("adios2"), nb::arg("mesh"), "Write mesh to file using ADIOS2"); +} + +template +void declare_write_meshtags(nb::module_& m) +{ + // dolfinx::io::native::write_meshtags + m.def( + "write_meshtags", + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, dolfinx::mesh::Mesh& mesh, + dolfinx::mesh::MeshTags& meshtags) + { + auto io = ADIOS2.io(); + auto engine = ADIOS2.engine(); + return dolfinx::io::native::write_meshtags(*io, *engine, mesh, + meshtags); + }, + nb::arg("adios2"), nb::arg("mesh"), nb::arg("meshtags"), + "Write meshtags to file using ADIOS2"); +} + +template +void declare_read_meshtags(nb::module_& m) +{ + // dolfinx::io::native::read_meshtags + m.def( + "read_meshtags", + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, dolfinx::mesh::Mesh& mesh) + { + auto io = ADIOS2.io(); + auto engine = ADIOS2.engine(); + + std::map> tags; + + for (unsigned int step = 0; + engine->BeginStep() == adios2::StepStatus::OK; ++step) + { + adios2::Attribute var_names + = io->InquireAttribute("meshtags_names"); + adios2::Attribute var_dtypes + = io->InquireAttribute("meshtags_dtypes"); + std::vector names = var_names.Data(); + std::vector dtypes = var_dtypes.Data(); + std::string name = names.back(); + std::string dtype = dtypes.back(); + + if (dtype == "int32_t") + { + dolfinx::mesh::MeshTags mt + = dolfinx::io::native::read_meshtags( + *io, *engine, mesh, name); + engine->EndStep(); + auto it = tags.end(); + tags.insert( + it, + std::pair>( + name, mt)); + } + else + { + throw std::runtime_error( + "The datatype associated with the meshtags values " + "is not supported yet"); + } + } + return tags; + }, + nb::arg("adios2"), nb::arg("mesh"), + "Read all meshtags from file using ADIOS2"); +} + +#endif + } // namespace void io(nb::module_& m) { +#ifdef HAS_ADIOS2 + // dolfinx::io::ADIOS2Wrapper + nb::class_ ADIOS2(m, "ADIOS2"); + + ADIOS2 + .def( + "__init__", + [](dolfinx::io::ADIOS2Wrapper* v, MPICommWrapper comm, + const std::string filename, std::string tag, + std::string engine_type = "BP5", std::string mode = "write") + { + new (v) dolfinx::io::ADIOS2Wrapper(comm.get(), filename, tag, + engine_type, mode); + }, + nb::arg("comm"), nb::arg("filename"), nb::arg("tag"), + nb::arg("engine_type"), nb::arg("mode")) + .def( + "__init__", + [](dolfinx::io::ADIOS2Wrapper* v, std::string config, + MPICommWrapper comm, const std::string filename, std::string tag, + std::string mode = "write") + { + new (v) dolfinx::io::ADIOS2Wrapper(config, comm.get(), filename, + tag, mode); + }, + nb::arg("config"), nb::arg("comm"), nb::arg("filename"), + nb::arg("tag"), nb::arg("mode")) + .def("close", &dolfinx::io::ADIOS2Wrapper::close); + + // dolfinx::io::impl_native::read_mesh_variant + m.def( + "read_mesh", + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, MPICommWrapper comm, + dolfinx::mesh::GhostMode ghost_mode) + { + auto io = ADIOS2.io(); + auto engine = ADIOS2.engine(); + return dolfinx::io::impl_native::read_mesh_variant( + *io, *engine, comm.get(), ghost_mode); + }, + nb::arg("adios2"), nb::arg("comm"), nb::arg("ghost_mode"), + "Read mesh from file using ADIOS2"); + + declare_write_mesh(m, "float32"); + declare_write_mesh(m, "float64"); + + // TODO: Include MeshTags of other types + declare_write_meshtags(m); + declare_write_meshtags(m); + declare_read_meshtags(m); + declare_read_meshtags(m); + +#endif + // dolfinx::io::cell vtk cell type converter m.def("get_vtk_cell_type", &dolfinx::io::cells::get_vtk_cell_type, nb::arg("cell"), nb::arg("dim"), "Get VTK cell identifier"); diff --git a/python/test/unit/io/test_adios2.py b/python/test/unit/io/test_adios2.py index 980558415dd..4d4e7a2c967 100644 --- a/python/test/unit/io/test_adios2.py +++ b/python/test/unit/io/test_adios2.py @@ -4,19 +4,32 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later +import typing +from collections import ChainMap from pathlib import Path from mpi4py import MPI import numpy as np +import numpy.typing as npt import pytest import ufl from basix.ufl import element from dolfinx import default_real_type, default_scalar_type -from dolfinx.fem import Function, functionspace +from dolfinx.fem import Function, assemble_scalar, form, functionspace from dolfinx.graph import adjacencylist -from dolfinx.mesh import CellType, create_mesh, create_unit_cube, create_unit_square +from dolfinx.mesh import ( + CellType, + GhostMode, + Mesh, + MeshTags, + compute_midpoints, + create_mesh, + create_unit_cube, + create_unit_square, + meshtags, +) def generate_mesh(dim: int, simplex: bool, N: int = 5, dtype=None): @@ -38,6 +51,173 @@ def generate_mesh(dim: int, simplex: bool, N: int = 5, dtype=None): raise RuntimeError("Unsupported dimension") +def generate_map( + mesh: Mesh, + meshtag: MeshTags, + comm: MPI.Intracomm, + root: int, +) -> typing.Optional[dict[str, tuple[int, npt.NDArray]]]: + """ + NOTE: From ADIOS4DOLFINx + Helper function to generate map from meshtag value to its corresponding index and midpoint. + + Args: + mesh: The mesh + meshtag: The associated meshtag + comm: MPI communicator to gather the map from all processes with + root (int): Rank to store data on + Returns: + Root rank returns the map, all other ranks return None + """ + mesh.topology.create_connectivity(meshtag.dim, mesh.topology.dim) + midpoints = compute_midpoints(mesh, meshtag.dim, meshtag.indices) + e_map = mesh.topology.index_map(meshtag.dim) + value_to_midpoint = {} + for index, value in zip(meshtag.indices, meshtag.values): + value_to_midpoint[value] = ( + int(e_map.local_range[0] + index), + midpoints[index], + ) + global_map = comm.gather(value_to_midpoint, root=root) + if comm.rank == root: + return dict(ChainMap(*global_map)) # type: ignore + return None + + +# TODO: Fix problems with ("HDF5", ".h5") +@pytest.mark.adios2 +@pytest.mark.parametrize("encoder, suffix", [("BP4", ".bp"), ("BP5", ".bp")]) +@pytest.mark.parametrize("ghost_mode", [GhostMode.shared_facet, GhostMode.none]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("simplex", [True, False]) +def test_mesh_read_write(encoder, suffix, ghost_mode, dtype, dim, simplex, tmp_path): + "Test writing of a mesh" + from dolfinx.io import ADIOS2, read_mesh, write_mesh + + N = 5 + # Consistent tmp dir across processes + fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) + file = fname / f"adios_mesh_{encoder}" + + mesh = generate_mesh(dim, simplex, N, dtype) + + adios = ADIOS2( + mesh.comm, + filename=str(file.with_suffix(suffix)), + tag="mesh-write", + engine_type=encoder, + mode="write", + ) + + write_mesh(adios, mesh) + + adios_read = ADIOS2( + MPI.COMM_WORLD, + filename=str(file.with_suffix(suffix)), + tag="mesh-read", + engine_type=encoder, + mode="read", + ) + + mesh_adios = read_mesh(adios_read, MPI.COMM_WORLD, ghost_mode=ghost_mode) + + mesh_adios.comm.Barrier() + mesh.comm.Barrier() + + for i in range(mesh.topology.dim + 1): + mesh.topology.create_entities(i) + mesh_adios.topology.create_entities(i) + assert ( + mesh.topology.index_map(i).size_global == mesh_adios.topology.index_map(i).size_global + ) + + # Check that integration over different entities are consistent + measures = [ufl.ds, ufl.dx] if ghost_mode is GhostMode.none else [ufl.ds, ufl.dS, ufl.dx] + for measure in measures: + c_adios = assemble_scalar(form(1 * measure(domain=mesh_adios), dtype=dtype)) + c_ref = assemble_scalar(form(1 * measure(domain=mesh), dtype=dtype)) + assert np.isclose( + mesh_adios.comm.allreduce(c_adios, MPI.SUM), + mesh.comm.allreduce(c_ref, MPI.SUM), + ) + + +# FIXME: ("BP4", ".bp") unable to read the shape of values +@pytest.mark.adios2 +@pytest.mark.parametrize("encoder, suffix", [("BP5", ".bp")]) +@pytest.mark.parametrize("ghost_mode", [GhostMode.shared_facet, GhostMode.none]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("simplex", [True, False]) +def test_meshtags(encoder, suffix, ghost_mode, dtype, dim, simplex, tmp_path): + """Test checkpointing of meshtags""" + from dolfinx.io import ADIOS2, read_mesh, read_meshtags, write_mesh, write_meshtags + + root = 0 + N = 5 + mesh = generate_mesh(dim, simplex, N, dtype) + + # Consistent tmp dir across processes + fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) + key = f"meshtags_{mesh.topology.cell_name()}_{encoder}" + file = fname / key + + adios = ADIOS2( + mesh.comm, + filename=str(file.with_suffix(suffix)), + tag="meshtags-write", + engine_type=encoder, + mode="write", + ) + + write_mesh(adios, mesh) + + # Create and write meshtags and store midpoints of entities + ref_maps = [] + for dim in range(mesh.topology.dim + 1): + mesh.topology.create_connectivity(dim, mesh.topology.dim) + e_map = mesh.topology.index_map(dim) + num_entities_local = e_map.size_local + entities = np.arange(num_entities_local, dtype=np.int32) + mt = meshtags(mesh, dim, entities, e_map.local_range[0] + entities) + mt.name = f"entity_{dim}" + + write_meshtags(adios, mesh, mt) + ref_map = generate_map(mesh, mt, mesh.comm, root) + ref_maps.append(ref_map) + + del mt + + del mesh + + adios.close() + MPI.COMM_WORLD.Barrier() + + adios_read = ADIOS2( + MPI.COMM_WORLD, + filename=str(file.with_suffix(suffix)), + tag="meshtags-read", + engine_type=encoder, + mode="read", + ) + + mesh_adios = read_mesh(adios_read, MPI.COMM_WORLD, ghost_mode=ghost_mode) + tags_read = read_meshtags(adios_read, mesh_adios) + + for dim, (mt_name, mt_adios) in enumerate(tags_read.items()): + assert dim == mt_adios.dim + + read_map = generate_map(mesh_adios, mt_adios, mesh_adios.comm, root) + + if MPI.COMM_WORLD.rank == root: + ref_map = ref_maps[dim] + assert len(ref_map) == len(read_map) + for value, (_, midpoint) in ref_map.items(): + _, read_midpoint = read_map[value] + np.testing.assert_allclose(read_midpoint, midpoint) + + @pytest.mark.adios2 class TestFides: @pytest.mark.parametrize("dim", [2, 3])