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..dd04dafe688 --- /dev/null +++ b/cpp/demo/checkpointing/main.cpp @@ -0,0 +1,133 @@ +// ```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; + +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)); + + 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); + std::span x = mesh->geometry().x(); + std::ranges::transform(x, x.begin(), [](auto xi) { return xi *= 4; }); + + io::native::write_mesh(io, engine, *mesh, 0.5); + + engine.Close(); + } + catch (std::exception& e) + { + std::cout << "ERROR: ADIOS2 exception: " << e.what() << "\n"; + MPI_Abort(MPI_COMM_WORLD, -1); + } + + try + { + // Read mode : 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); + + // ReadRandomAccess mode : set up ADIOS2 IO and Engine + adios2::IO io_rra = adios_read.DeclareIO("mesh-rra"); + io_rra.SetEngine("BP5"); + adios2::Engine engine_rra + = io_rra.Open("mesh.bp", adios2::Mode::ReadRandomAccess); + + // Write mode : set up ADIOS2 IO and Engine + 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); + + // Find the time stamps array + auto var_time = io_rra.InquireVariable("time"); + const std::vector::Info>> timestepsinfo + = var_time.AllStepsBlocksInfo(); + + std::size_t num_steps = timestepsinfo.size(); + std::vector times(num_steps); + + for (std::size_t step = 0; step < num_steps; ++step) + { + var_time.SetStepSelection({step, 1}); + engine_rra.Get(var_time, times[step]); + } + engine_rra.Close(); + + // Read mesh + engine_read.BeginStep(); + auto mesh_read + = io::native::read_mesh(io_read, engine_read, MPI_COMM_WORLD); + if (engine_read.BetweenStepPairs()) + { + engine_read.EndStep(); + } + // Write mesh + io::native::write_mesh(io_write, engine_write, mesh_read); + + // Update mesh + double time = 0.5; + std::size_t querystep; + auto pos = std::ranges::find(times, time); + if (pos != times.end()) + { + querystep = std::ranges::distance(times.begin(), pos); + std::cout << "Query step is : " << querystep << "\n"; + } + else + { + throw std::runtime_error("Step corresponding to time : " + + std::to_string(time) + " not found"); + } + + io::native::update_mesh(io_read, engine_read, mesh_read, querystep); + + // Write updated mesh + io::native::write_mesh(io_write, engine_write, mesh_read, time); + + engine_read.Close(); + 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..c5a0fb27c37 --- /dev/null +++ b/cpp/dolfinx/io/ADIOS2_utils.h @@ -0,0 +1,136 @@ +// 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}, + {"readrandomaccess", adios2::Mode::ReadRandomAccess}, +}; +} // 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 + ADIOS2Wrapper(MPI_Comm comm) + { + _adios = std::make_shared(comm); + } + + /// @brief Create an ADIOS2-based engine writer/reader + /// @param[in] config_file Path to config file to set up ADIOS2 engines from + /// @param[in] comm The MPI communicator + ADIOS2Wrapper(std::string config_file, MPI_Comm comm) + { + _adios = std::make_shared(config_file, comm); + } + + /// @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 Add IO and an Engine with specified type and mode + /// @param[in] filename Name of input/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, available are: "write", "read", "append", + /// "readrandomaccess", and the default is "append" + void add_io(const std::string filename, std::string tag, + std::string engine_type = "BP5", std::string mode = "append") + { + std::map>::iterator it_io + = _ios.end(); + _ios.insert(it_io, + std::pair>( + tag, std::make_shared(_adios->DeclareIO(tag)))); + + assert(_ios[tag]); + if (!_ios[tag]->InConfigFile()) + { + _ios[tag]->SetEngine(engine_type); + } + std::map>::iterator it_engine + = _engines.end(); + _engines.insert(it_engine, + std::pair>( + tag, std::make_shared(_ios[tag]->Open( + filename, string_to_mode[mode])))); + } + + /// @brief Close engine associated with the IO with the given tag + /// @param[in] tag The ADIOS2 IO name whose associated engine needs to be + /// closed + void close(std::string tag) + { + assert(_engines[tag]); + if (*_engines[tag]) + _engines[tag]->Close(); + } + + /// @brief Close all Engines + void close() + { + for (auto it = _engines.begin(); it != _engines.end(); ++it) + { + auto engine = it->second; + assert(engine); + if (*engine) + engine->Close(); + } + } + + /// @brief Get the IO with the given tag + /// @param[in] tag The ADIOS2 IO name + std::shared_ptr io(std::string tag) { return _ios[tag]; } + + /// @brief Get the Engine object + /// @param[in] tag The ADIOS2 IO name + std::shared_ptr engine(std::string tag) + { + return _engines[tag]; + } + +private: + std::shared_ptr _adios; + std::map> _ios; + std::map> _engines; +}; + +} // 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..326191608ea --- /dev/null +++ b/cpp/dolfinx/io/checkpointing.cpp @@ -0,0 +1,562 @@ +// 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}, +}; + +template +adios2::Variable define_var(adios2::IO& io, std::string name, + const adios2::Dims& shape = adios2::Dims(), + const adios2::Dims& start = adios2::Dims(), + const adios2::Dims& count = adios2::Dims()) +{ + if (adios2::Variable var = io.InquireVariable(name); var) + { + if (var.Count() != count) + var.SetSelection({start, count}); + + return var; + } + else + return io.DefineVariable(name, shape, start, count, + adios2::ConstantDims); +} + +} // namespace + +namespace dolfinx::io::native +{ +//----------------------------------------------------------------------------- +template +void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh, double time) +{ + + const mesh::Geometry& geometry = mesh.geometry(); + std::shared_ptr topology = mesh.topology(); + assert(topology); + std::shared_ptr geom_imap = geometry.index_map(); + + std::size_t currentstep = engine.CurrentStep(); + assert(!engine.BetweenStepPairs()); + engine.BeginStep(); + + if (currentstep == 0) + { + // Variables/attributes to save + std::int32_t dim = geometry.dim(); + std::int32_t tdim = topology->dim(); + std::uint64_t num_cells_global, cell_offset; + std::uint32_t num_cells_local, degree; + std::string cell_type, lagrange_variant; + std::vector array_global; + std::vector offsets_global; + + // 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()]; + } + + 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_cells + = io.DefineVariable("num_cells"); + + 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); + + engine.Put(var_num_cells, num_cells_global); + engine.Put(var_topology_array, array_global.data()); + engine.Put(var_topology_offsets, offsets_global.data()); + } + } + + std::uint64_t num_nodes_global, offset; + std::uint32_t num_nodes_local; + // 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]; + } + const std::span mesh_x = geometry.x(); + + if (currentstep == 0) + { + adios2::Variable var_num_nodes + = io.DefineVariable("num_nodes"); + engine.Put(var_num_nodes, num_nodes_global); + } + + adios2::Variable var_time = define_var(io, "time", {}, {}, {}); + + adios2::Variable var_x = define_var(io, "x", {num_nodes_global, 3}, + {offset, 0}, {num_nodes_local, 3}); + + engine.Put(var_time, time); + engine.Put(var_x, mesh_x.subspan(0, num_nodes_local * 3).data()); + + engine.EndStep(); + + spdlog::info("Mesh written"); +} + +//----------------------------------------------------------------------------- +/// @cond +template void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh, + double time); + +template void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh, + double time); + +/// @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 + +//----------------------------------------------------------------------------- +template +void update_mesh(adios2::IO& io, adios2::Engine& engine, + dolfinx::mesh::Mesh& mesh, std::size_t step) +{ + if (!engine.BetweenStepPairs()) + { + engine.BeginStep(); + } + + std::uint32_t num_nodes_local = mesh.geometry().index_map()->size_local(); + std::vector x_raw(num_nodes_local * 3); + + // Read variables + { + double time; + adios2::Variable var_time = io.InquireVariable("time"); + var_time.SetStepSelection({step, 1}); + if (var_time) + { + engine.Get(var_time, time); + spdlog::info("Updating geometry at time : {}", time); + } + else + { + throw std::runtime_error("Step : " + std::to_string(step) + " not found"); + } + + adios2::Variable var_x = io.InquireVariable("x"); + var_x.SetStepSelection({step, 1}); + if (var_x) + { + std::uint64_t nodes_offset + = mesh.geometry().index_map()->local_range()[0]; + var_x.SetSelection({{nodes_offset, 0}, {num_nodes_local, 3}}); + engine.Get(var_x, x_raw.data(), adios2::Mode::Sync); + } + else + { + throw std::runtime_error("Coordinates data not found at step : " + step); + } + } + + engine.EndStep(); + + // Redistribute adios2 input coordinate data and find updated coordinates of + // the mesh + std::vector x_new = dolfinx::MPI::distribute_data( + mesh.comm(), mesh.geometry().input_global_indices(), mesh.comm(), x_raw, + 3); + + std::span x = mesh.geometry().x(); + std::ranges::transform(x_new, x.begin(), [](auto it) { return it; }); +} + +//----------------------------------------------------------------------------- +/// @cond +template void update_mesh(adios2::IO& io, adios2::Engine& engine, + dolfinx::mesh::Mesh& mesh, + std::size_t step); + +template void update_mesh(adios2::IO& io, adios2::Engine& engine, + dolfinx::mesh::Mesh& mesh, + std::size_t step); + +/// @endcond + +//----------------------------------------------------------------------------- +std::vector read_timestamps(adios2::IO& io, adios2::Engine& engine) +{ + if (engine.OpenMode() != adios2::Mode::ReadRandomAccess) + { + throw std::runtime_error( + "Time stamps can only be read in ReadRandomAccess mode"); + } + adios2::Variable var_time = io.InquireVariable("time"); + const std::vector::Info>> timestepsinfo + = var_time.AllStepsBlocksInfo(); + + std::size_t num_steps = timestepsinfo.size(); + std::vector times(num_steps); + + for (std::size_t step = 0; step < num_steps; ++step) + { + var_time.SetStepSelection({step, 1}); + engine.Get(var_time, times[step]); + } + return times; +} + +} // 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..8fe1579efcf --- /dev/null +++ b/cpp/dolfinx/io/checkpointing.h @@ -0,0 +1,121 @@ +// 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 +#include +#include +#include +#include + +/// @file checkpointing.h +/// @brief ADIOS2 based checkpointing + +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 +/// @param[in] time Time associated with the mesh with moving geometry +template +void write_mesh(adios2::IO& io, adios2::Engine& engine, + const dolfinx::mesh::Mesh& mesh, double time = 0); + +/// @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 Update geometry of mesh by reading from a file. +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @param[in] mesh Mesh of type float or double +/// @param[in] step ADIOS2 Engine Step to read geometry from +/// @note This method expects that ReadRandomAccess mode is used to determine +/// the step in which a given time stamp exists +template +void update_mesh(adios2::IO& io, adios2::Engine& engine, + dolfinx::mesh::Mesh& mesh, std::size_t step); + +/// @brief Read time stamps. +/// +/// @param[in] io ADIOS2 IO +/// @param[in] engine ADIOS2 Engine +/// @note The Engine should be opened in ReadRandomAccess mode +std::vector read_timestamps(adios2::IO& io, adios2::Engine& engine); + +} // 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..5de8075db38 --- /dev/null +++ b/python/demo/demo_checkpointing.py @@ -0,0 +1,95 @@ +# --- +# 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 numpy as np + +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 + +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, +) +# - + +# + +config_path = os.getcwd() + "/checkpointing.yml" +adios2 = io.ADIOS2(config_path, msh.comm) +tag = "mesh-write" +adios2.add_io(filename="mesh.bp", tag=tag, mode="write") +# - + +# + +io.write_mesh(adios2, tag, msh) + +msh.geometry.x[:] += 4 +io.write_mesh(adios2, tag, msh, 0.5) + +msh.geometry.x[:] += 4 +io.write_mesh(adios2, tag, msh, 1.0) + +adios2.close(tag) +# - + +# + +tag = "mesh-readrandomaccess" +adios2.add_io(filename="mesh.bp", tag=tag, engine_type="BP5", mode="readrandomaccess") +# - + +# + +times = io.read_timestamps(adios2, tag) +print(f"Time stamps : {times}") +# - + +# + +tag = "mesh-read" +adios2.add_io(filename="mesh.bp", tag=tag, engine_type="BP5", mode="read") +# - + +# + +msh_read = io.read_mesh(adios2, tag, msh.comm) +print(np.max(msh_read.geometry.x)) + +io.update_mesh(adios2, tag, msh_read, 1) +print(np.max(msh_read.geometry.x)) + +io.update_mesh(adios2, tag, msh_read, 2) +print(np.max(msh_read.geometry.x)) + +adios2.close(tag) +# - 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/io/__init__.py b/python/dolfinx/io/__init__.py index 3a53be8d7eb..cd6eff9e1ea 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_timestamps, + update_mesh, + write_mesh, + ) - __all__ = [*__all__, "FidesWriter", "VTXWriter", "FidesMeshPolicy", "VTXMeshPolicy"] + __all__ = [ + *__all__, + "FidesWriter", + "VTXWriter", + "FidesMeshPolicy", + "VTXMeshPolicy", + "ADIOS2", + "read_mesh", + "read_timestamps", + "update_mesh", + "write_mesh", + ] diff --git a/python/dolfinx/io/utils.py b/python/dolfinx/io/utils.py index a479f898998..445690dea73 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_timestamps", + "update_mesh", + "write_mesh", + ] class VTXWriter: """Writer for VTX files, using ADIOS2 to create the files. @@ -186,6 +203,100 @@ def write(self, t: float): def close(self): self._cpp_object.close() + class ADIOS2(_cpp.io.ADIOS2): + """Wrapper around ADIOS2. + Supports creation of arbitrary number of IOs + """ + + def __enter__(self): + return self + + def __exit__(self, exception_type, exception_value, traceback): + self.close() + + def add_io(self, filename: str, tag: str, engine_type: str = "BP5", mode: str = "append"): + """Add IO and Engine + + Args: + filename: filename to associate Engine with + tag: Name of the IO to use + engine_type: ADIOS2 Engine type, default is "BP5" + mode: ADIOS2 mode, default is "append" + """ + super().add_io(filename, tag, engine_type, mode) + + def close(self, tag: str): + """Close IO and Engine associated with the given tag""" + super().close(tag) + + def write_mesh(ADIOS2: ADIOS2, tag: str, mesh: Mesh, time: np.floating = 0) -> None: + """Write mesh to a file using ADIOS2 + + Args: + ADIOS2: Wrapper around ADIOS2. + tag: Name of the IO to use. + mesh: Mesh to write to the file. + time: Associated time stamp + """ + + 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, tag, mesh._cpp_object, time) + + def read_mesh( + ADIOS2: ADIOS2, tag: str, comm: _MPI.Comm, ghost_mode: GhostMode = GhostMode.shared_facet + ) -> Mesh: + """Read mesh from a file using ADIOS2 + + Args: + ADIOS2: Wrapper around ADIOS2 + tag: Name of the IO to use + comm: The MPI communicator + ghost_mode: GhostMode to use on mesh partitioning + Returns: + Mesh + """ + msh = _cpp.io.read_mesh(ADIOS2, tag, 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 read_timestamps(ADIOS2: ADIOS2, tag: str) -> np.ndarray: + """Read timestamps from a file using ADIOS2 + + Args: + ADIOS2: Wrapper around ADIOS2 + tag: Name of the IO to use + Returns: + vector of timestamps + """ + return _cpp.io.read_timestamps(ADIOS2, tag) + + def update_mesh(ADIOS2: ADIOS2, tag: str, mesh: Mesh, step: np.int64) -> None: + """Update mesh geometry with geometry stored with the given time stamp in the file + + Args: + ADIOS2: Wrapper around ADIOS2 + tag: Name of the IO to use + mesh: Mesh to update the geometry + step: ADIOS2 Step at which the corresponding geometry is to be read from the file + """ + + return _cpp.io.update_mesh(ADIOS2, tag, mesh._cpp_object, step) + 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..c10deb05506 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 @@ -232,10 +234,111 @@ 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, std::string tag, + dolfinx::mesh::Mesh& mesh, double time) + { + auto io = ADIOS2.io(tag); + auto engine = ADIOS2.engine(tag); + return dolfinx::io::native::write_mesh(*io, *engine, mesh, time); + }, + nb::arg("adios2"), nb::arg("tag"), nb::arg("mesh"), nb::arg("time") = 0.0, + "Write mesh to file using ADIOS2"); +} + +template +void declare_update_mesh(nb::module_& m) +{ + // dolfinx::io::native::update_mesh + m.def( + "update_mesh", + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, std::string tag, + dolfinx::mesh::Mesh& mesh, std::size_t step) + { + auto io = ADIOS2.io(tag); + auto engine = ADIOS2.engine(tag); + return dolfinx::io::native::update_mesh(*io, *engine, mesh, step); + }, + nb::arg("adios2"), nb::arg("tag"), nb::arg("mesh"), nb::arg("step"), + "Update mesh with geometry associated with a given ADIOS2 step"); +} + +#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) + { new (v) dolfinx::io::ADIOS2Wrapper(comm.get()); }, nb::arg("comm")) + .def( + "__init__", + [](dolfinx::io::ADIOS2Wrapper* v, std::string config, + MPICommWrapper comm) + { new (v) dolfinx::io::ADIOS2Wrapper(config, comm.get()); }, + nb::arg("config"), nb::arg("comm")) + .def( + "add_io", + [](dolfinx::io::ADIOS2Wrapper& self, const std::string filename, + std::string tag, std::string engine_type = "BP5", + std::string mode = "append") + { self.add_io(filename, tag, engine_type, mode); }, + nb::arg("filename"), nb::arg("tag"), nb::arg("engine_type"), + nb::arg("mode"), "Create IO and Engine") + .def( + "close", [](dolfinx::io::ADIOS2Wrapper& self) { self.close(); }, + "Close all engines") + .def( + "close", [](dolfinx::io::ADIOS2Wrapper& self, std::string tag) + { self.close(tag); }, nb::arg("tag"), + "Close engine associated with tag"); + + // dolfinx::io::impl_native::read_mesh_variant + m.def( + "read_mesh", + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, std::string tag, + MPICommWrapper comm, dolfinx::mesh::GhostMode ghost_mode) + { + auto io = ADIOS2.io(tag); + auto engine = ADIOS2.engine(tag); + return dolfinx::io::impl_native::read_mesh_variant( + *io, *engine, comm.get(), ghost_mode); + }, + nb::arg("adios2"), nb::arg("tag"), nb::arg("comm"), nb::arg("ghost_mode"), + "Read mesh from file using ADIOS2"); + + // dolfinx::io::native::read_timestamps + m.def( + "read_timestamps", + [](dolfinx::io::ADIOS2Wrapper& ADIOS2, std::string tag) + { + auto io = ADIOS2.io(tag); + auto engine = ADIOS2.engine(tag); + return dolfinx::io::native::read_timestamps(*io, *engine); + }, + nb::arg("adios2"), nb::arg("tag"), + "Update mesh with geometry associated with a given ADIOS2 step"); + + declare_write_mesh(m, "float32"); + declare_write_mesh(m, "float64"); + declare_update_mesh(m); + declare_update_mesh(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 cad2c2a697e..8dfe10443ae 100644 --- a/python/test/unit/io/test_adios2.py +++ b/python/test/unit/io/test_adios2.py @@ -14,9 +14,9 @@ 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, create_mesh, create_unit_cube, create_unit_square def generate_mesh(dim: int, simplex: bool, N: int = 5, dtype=None): @@ -38,6 +38,154 @@ def generate_mesh(dim: int, simplex: bool, N: int = 5, dtype=None): raise RuntimeError("Unsupported dimension") +# TODO: Fix problems with ("HDF5", ".h5"), ("BP4", ".bp"), +@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_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) + tag = "mesh-write" + adios.add_io(filename=str(file.with_suffix(suffix)), tag=tag, engine_type=encoder, mode="write") + + write_mesh(adios, tag, mesh) + + adios_read = ADIOS2(MPI.COMM_WORLD) + tag = "mesh-read" + adios_read.add_io( + filename=str(file.with_suffix(suffix)), tag=tag, engine_type=encoder, mode="read" + ) + + mesh_adios = read_mesh(adios_read, tag, 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), + ) + + +# TODO: Fix problems with ("HDF5", ".h5"), ("BP4", ".bp"), +@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_timedep_mesh_read_write(encoder, suffix, ghost_mode, dtype, dim, simplex, tmp_path): + "Test writing of a time dependent mesh" + from dolfinx.io import ADIOS2, read_mesh, read_timestamps, update_mesh, write_mesh + + N = 5 + # Consistent tmp dir across processes + fname = MPI.COMM_WORLD.bcast(tmp_path, root=0) + file = fname / f"adios_timedep_mesh_{encoder}" + + mesh = generate_mesh(dim, simplex, N, dtype) + + def displace(x): + return np.asarray( + [ + x[0] + 0.1 * np.sin(x[0]) * np.cos(x[1]), + x[1] + 0.6 * np.cos(x[0]) * np.sin(x[1]), + x[2], + ] + ) + + adios = ADIOS2(mesh.comm) + tag_write = "mesh-write" + adios.add_io( + filename=str(file.with_suffix(suffix)), tag=tag_write, engine_type=encoder, mode="write" + ) + + # Write mesh + write_mesh(adios, tag_write, mesh, time=0.0) + + delta_x1 = displace(mesh.geometry.x.T).T + mesh.geometry.x[:] += delta_x1 + + write_mesh(adios, tag_write, mesh, time=1.0) + + delta_x2 = displace(mesh.geometry.x.T).T + mesh.geometry.x[:] += delta_x2 + + write_mesh(adios, tag_write, mesh, time=2.0) + + adios.close(tag_write) + + # reset mesh geometry to the one at time=0.0 + mesh.geometry.x[:] -= delta_x1 + delta_x2 + + tag_rra = "mesh-readrandomaccess" + adios.add_io( + filename=str(file.with_suffix(suffix)), + tag=tag_rra, + engine_type=encoder, + mode="readrandomaccess", + ) + times = read_timestamps(adios, tag_rra) + adios.close(tag_rra) + assert np.all(np.isclose(times, [0.0, 1.0, 2.0])) + + tag_read = "mesh-read" + adios.add_io( + filename=str(file.with_suffix(suffix)), tag=tag_read, engine_type=encoder, mode="read" + ) + mesh_adios = read_mesh(adios, tag_read, MPI.COMM_WORLD, ghost_mode=ghost_mode) + + mesh_adios.comm.Barrier() + mesh.comm.Barrier() + + # 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 step, time in enumerate(times): + if step == 1: + mesh.geometry.x[:] += delta_x1 + if step == 2: + mesh.geometry.x[:] += delta_x2 + + # FIXME: update_mesh at time time=0.0 should work!? + if step > 0: + update_mesh(adios, tag_read, mesh_adios, step) + + mesh_adios.comm.Barrier() + mesh.comm.Barrier() + + 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), + ) + + @pytest.mark.adios2 class TestFides: @pytest.mark.parametrize("dim", [2, 3])