VTKHDF mesh I/O (#3576)

* Work in progress

* Work in progress

* Reorganise

* tidy

* Add comments and test

* Interface update

* Work on Python layer

* Add docstrings

* Add mixed mesh fixture

* fix?

* Add typing

---------

Co-authored-by: Garth N. Wells <[email protected]>

cpp/dolfinx/io/CMakeLists.txt

@@ -1,26 +1,27 @@
 set(HEADERS_io
-    ${CMAKE_CURRENT_SOURCE_DIR}/dolfinx_io.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/cells.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/XDMFFile.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_function.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_mesh.h
-    ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_utils.h
-    PARENT_SCOPE
+  ${CMAKE_CURRENT_SOURCE_DIR}/dolfinx_io.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/cells.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/VTKHDF.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/XDMFFile.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_function.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_mesh.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_utils.h
+  PARENT_SCOPE
 )
 
 target_sources(
   dolfinx
   PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/ADIOS2Writers.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/cells.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/XDMFFile.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_function.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_mesh.cpp
-          ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_utils.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/cells.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/HDF5Interface.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/VTKFile.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/vtk_utils.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/XDMFFile.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_function.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_mesh.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/xdmf_utils.cpp
 )

cpp/dolfinx/io/HDF5Interface.h

@@ -37,6 +37,8 @@ hid_t hdf5_type()
     return H5T_NATIVE_INT64;
   else if constexpr (std::is_same_v<T, std::uint64_t>)
     return H5T_NATIVE_UINT64;
+  else if constexpr (std::is_same_v<T, std::uint8_t>)
+    return H5T_NATIVE_UINT8;
   else if constexpr (std::is_same_v<T, std::size_t>)
   {
     throw std::runtime_error(

cpp/dolfinx/io/VTKHDF.h

@@ -0,0 +1,301 @@
+// Copyright (C) 2024 Chris Richardson
+//
+// This file is part of DOLFINx (https://www.fenicsproject.org)
+//
+// SPDX-License-Identifier:    LGPL-3.0-or-later
+
+#include "HDF5Interface.h"
+
+#include <dolfinx/common/IndexMap.h>
+#include <dolfinx/io/cells.h>
+#include <dolfinx/mesh/Mesh.h>
+#include <dolfinx/mesh/Topology.h>
+#include <dolfinx/mesh/utils.h>
+
+#include <map>
+#include <vector>
+
+namespace dolfinx::io::VTKHDF
+{
+
+/// @brief Write a mesh to VTKHDF format
+/// @tparam U Scalar type of the mesh
+/// @param filename File to write to
+/// @param mesh Mesh
+template <typename U>
+void write_mesh(std::string filename, const mesh::Mesh<U>& mesh);
+
+/// @brief  Read a mesh from VTKHDF format file
+/// @tparam U Scalar type of mesh
+/// @param comm MPI Communicator for reading mesh
+/// @param filename Filename
+/// @return a mesh
+template <typename U>
+mesh::Mesh<U> read_mesh(MPI_Comm comm, std::string filename);
+
+} // namespace dolfinx::io::VTKHDF
+
+using namespace dolfinx;
+
+template <typename U>
+void io::VTKHDF::write_mesh(std::string filename, const mesh::Mesh<U>& mesh)
+{
+  hid_t h5file = io::hdf5::open_file(mesh.comm(), filename, "w", true);
+
+  // Create VTKHDF group
+  io::hdf5::add_group(h5file, "VTKHDF");
+
+  hid_t vtk_group = H5Gopen(h5file, "VTKHDF", H5P_DEFAULT);
+
+  // Create "Version" attribute
+  hsize_t dims = 2;
+  hid_t space_id = H5Screate_simple(1, &dims, NULL);
+  hid_t attr_id = H5Acreate(vtk_group, "Version", H5T_NATIVE_INT32, space_id,
+                            H5P_DEFAULT, H5P_DEFAULT);
+  std::array<std::int32_t, 2> version = {2, 2};
+  H5Awrite(attr_id, H5T_NATIVE_INT32, version.data());
+  H5Aclose(attr_id);
+  H5Sclose(space_id);
+
+  // Create "Type" attribute
+  space_id = H5Screate(H5S_SCALAR);
+  hid_t atype = H5Tcopy(H5T_C_S1);
+  H5Tset_size(atype, 16);
+  H5Tset_strpad(atype, H5T_STR_NULLTERM);
+  attr_id
+      = H5Acreate(vtk_group, "Type", atype, space_id, H5P_DEFAULT, H5P_DEFAULT);
+  H5Awrite(attr_id, atype, "UnstructuredGrid");
+
+  H5Aclose(attr_id);
+  H5Sclose(space_id);
+  H5Gclose(vtk_group);
+
+  // Extract topology information for each cell type
+  auto cell_types = mesh.topology()->entity_types(mesh.topology()->dim());
+
+  std::vector cell_index_maps
+      = mesh.topology()->index_maps(mesh.topology()->dim());
+  std::vector<std::int32_t> num_cells;
+  std::vector<std::int64_t> num_cells_global;
+  for (auto im : cell_index_maps)
+  {
+    num_cells.push_back(im->size_local());
+    num_cells_global.push_back(im->size_global());
+  }
+
+  // Geometry dofmap and points
+  auto geom_imap = mesh.geometry().index_map();
+  std::int32_t gdim = mesh.geometry().dim();
+  std::int64_t size_global = geom_imap->size_global();
+  std::vector<std::int64_t> geom_global_shape = {size_global, gdim};
+  auto geom_irange = geom_imap->local_range();
+
+  io::hdf5::write_dataset(h5file, "/VTKHDF/Points", mesh.geometry().x().data(),
+                          geom_irange, geom_global_shape, true, false);
+
+  io::hdf5::write_dataset(h5file, "VTKHDF/NumberOfPoints", &size_global, {0, 1},
+                          {1}, true, false);
+
+  // VTKHDF stores the cells as an adjacency list,                                \
+  //        where cell types might be jumbled up.
+  std::vector<std::int64_t> topology_flattened;
+  std::vector<std::int64_t> topology_offsets;
+  std::vector<std::uint8_t> vtkcelltypes;
+  for (int i = 0; i < cell_index_maps.size(); ++i)
+  {
+    auto g_dofmap = mesh.geometry().dofmap(i);
+
+    std::vector<std::int32_t> local_dm;
+    local_dm.reserve(g_dofmap.extent(1) * num_cells[i]);
+    auto perm = io::cells::perm_vtk(cell_types[i], g_dofmap.extent(1));
+    for (int j = 0; j < num_cells[i]; ++j)
+      for (int k = 0; k < g_dofmap.extent(1); ++k)
+        local_dm.push_back(g_dofmap(j, perm[k]));
+
+    std::vector<std::int64_t> global_dm(local_dm.size());
+    geom_imap->local_to_global(local_dm, global_dm);
+
+    topology_flattened.insert(topology_flattened.end(), global_dm.begin(),
+                              global_dm.end());
+
+    topology_offsets.insert(topology_offsets.end(), g_dofmap.extent(0),
+                            g_dofmap.extent(1));
+
+    vtkcelltypes.insert(
+        vtkcelltypes.end(), cell_index_maps[i]->size_local(),
+        io::cells::get_vtk_cell_type(cell_types[i], mesh.topology()->dim()));
+  }
+  // Create topo_offsets
+  std::partial_sum(topology_offsets.cbegin(), topology_offsets.cend(),
+                   topology_offsets.begin());
+
+  std::vector<int> num_nodes_per_cell;
+  std::vector<std::int64_t> cell_start_pos;
+  std::vector<std::int64_t> cell_stop_pos;
+  for (int i = 0; i < cell_index_maps.size(); ++i)
+  {
+    num_nodes_per_cell.push_back(mesh.geometry().cmaps()[i].dim());
+    auto r = cell_index_maps[i]->local_range();
+    cell_start_pos.push_back(r[0]);
+    cell_stop_pos.push_back(r[1]);
+  }
+
+  // Compute overall cell offset from offsets for each cell type
+  std::int64_t offset_start_position
+      = std::accumulate(cell_start_pos.begin(), cell_start_pos.end(), 0);
+  std::int64_t offset_stop_position
+      = std::accumulate(cell_stop_pos.begin(), cell_stop_pos.end(), 0);
+
+  // Compute overall topology offset from offsets for each cell type
+  std::int64_t topology_start
+      = std::inner_product(num_nodes_per_cell.begin(), num_nodes_per_cell.end(),
+                           cell_start_pos.begin(), 0);
+
+  std::for_each(topology_offsets.begin(), topology_offsets.end(),
+                [topology_start](std::int64_t& t) { t += topology_start; });
+
+  std::int64_t num_all_cells_global
+      = std::accumulate(num_cells_global.begin(), num_cells_global.end(), 0);
+  io::hdf5::write_dataset(h5file, "/VTKHDF/Offsets", topology_offsets.data(),
+                          {offset_start_position + 1, offset_stop_position + 1},
+                          {num_all_cells_global + 1}, true, false);
+
+  // Store global mesh connectivity
+  std::int64_t topology_size_global
+      = std::inner_product(num_nodes_per_cell.begin(), num_nodes_per_cell.end(),
+                           num_cells_global.begin(), 0);
+
+  std::int64_t topology_stop = topology_start + topology_flattened.size();
+  io::hdf5::write_dataset(
+      h5file, "/VTKHDF/Connectivity", topology_flattened.data(),
+      {topology_start, topology_stop}, {topology_size_global}, true, false);
+
+  // Store cell types
+  io::hdf5::write_dataset(h5file, "/VTKHDF/Types", vtkcelltypes.data(),
+                          {offset_start_position, offset_stop_position},
+                          {num_all_cells_global}, true, false);
+
+  io::hdf5::write_dataset(h5file, "/VTKHDF/NumberOfConnectivityIds",
+                          &topology_size_global, {0, 1}, {1}, true, false);
+
+  io::hdf5::write_dataset(h5file, "/VTKHDF/NumberOfCells",
+                          &num_all_cells_global, {0, 1}, {1}, true, false);
+
+  io::hdf5::close_file(h5file);
+}
+
+template <typename U>
+mesh::Mesh<U> io::VTKHDF::read_mesh(MPI_Comm comm, std::string filename)
+{
+  hid_t h5file = io::hdf5::open_file(comm, filename, "r", true);
+
+  std::vector<std::int64_t> shape
+      = io::hdf5::get_dataset_shape(h5file, "/VTKHDF/Types");
+  int rank = MPI::rank(comm);
+  int size = MPI::size(comm);
+  auto local_cell_range = dolfinx::MPI::local_range(rank, shape[0], size);
+
+  hid_t dset_id = io::hdf5::open_dataset(h5file, "/VTKHDF/Types");
+  std::vector<std::uint8_t> types
+      = io::hdf5::read_dataset<std::uint8_t>(dset_id, local_cell_range, true);
+  H5Dclose(dset_id);
+  std::vector<std::uint8_t> types_unique(types.begin(), types.end());
+  {
+    std::ranges::sort(types_unique);
+    auto [unique_end, range_end] = std::ranges::unique(types_unique);
+    types_unique.erase(unique_end, range_end);
+  }
+
+  // Share cell types with all processes to make global list of cell types
+  // FIXME: amount of data is small, but number of connections does not scale.
+  std::int32_t count = types_unique.size();
+  std::vector<std::int32_t> recv_count(size);
+  std::vector<std::int32_t> recv_offsets(size + 1, 0);
+  MPI_Allgather(&count, 1, MPI_INT32_T, recv_count.data(), 1, MPI_INT32_T,
+                comm);
+  std::partial_sum(recv_count.begin(), recv_count.end(),
+                   recv_offsets.begin() + 1);
+  std::vector<std::uint8_t> recv_types(recv_offsets.back());
+  MPI_Allgatherv(types_unique.data(), count, MPI_UINT8_T, recv_types.data(),
+                 recv_count.data(), recv_offsets.data(), MPI_UINT8_T, comm);
+  {
+    std::ranges::sort(recv_types);
+    auto [unique_end, range_end] = std::ranges::unique(recv_types);
+    recv_types.erase(unique_end, range_end);
+  }
+
+  // Create reverse map from VTK to dolfinx cell types
+  std::map<std::uint8_t, mesh::CellType> vtk_to_dolfinx;
+  const std::vector<mesh::CellType> dolfinx_cells
+      = {mesh::CellType::point,       mesh::CellType::interval,
+         mesh::CellType::triangle,    mesh::CellType::quadrilateral,
+         mesh::CellType::tetrahedron, mesh::CellType::prism,
+         mesh::CellType::pyramid,     mesh::CellType::hexahedron};
+  for (auto dolfinx_type : dolfinx_cells)
+  {
+    vtk_to_dolfinx.insert({io::cells::get_vtk_cell_type(
+                               dolfinx_type, mesh::cell_dim(dolfinx_type)),
+                           dolfinx_type});
+  }
+
+  // Map from VTKCellType to index in list of cell types
+  std::map<std::uint8_t, std::int32_t> type_to_index;
+  std::vector<mesh::CellType> dolfinx_cell_type;
+  for (std::size_t i = 0; i < recv_types.size(); ++i)
+  {
+    type_to_index.insert({recv_types[i], i});
+    dolfinx_cell_type.push_back(vtk_to_dolfinx.at(recv_types[i]));
+  }
+
+  dset_id = io::hdf5::open_dataset(h5file, "/VTKHDF/NumberOfPoints");
+  std::vector<std::int64_t> npoints
+      = io::hdf5::read_dataset<std::int64_t>(dset_id, {0, 1}, true);
+  H5Dclose(dset_id);
+  spdlog::info("Mesh with {} points", npoints[0]);
+  auto local_point_range = MPI::local_range(rank, npoints[0], size);
+
+  std::vector<std::int64_t> x_shape
+      = io::hdf5::get_dataset_shape(h5file, "/VTKHDF/Points");
+  dset_id = io::hdf5::open_dataset(h5file, "/VTKHDF/Points");
+  std::vector<U> points_local
+      = io::hdf5::read_dataset<U>(dset_id, local_point_range, true);
+  H5Dclose(dset_id);
+  dset_id = io::hdf5::open_dataset(h5file, "/VTKHDF/Offsets");
+  std::vector<std::int64_t> offsets = io::hdf5::read_dataset<std::int64_t>(
+      dset_id, {local_cell_range[0], local_cell_range[1] + 1}, true);
+  H5Dclose(dset_id);
+  dset_id = io::hdf5::open_dataset(h5file, "/VTKHDF/Connectivity");
+  std::vector<std::int64_t> topology = io::hdf5::read_dataset<std::int64_t>(
+      dset_id, {offsets.front(), offsets.back()}, true);
+  H5Dclose(dset_id);
+  const std::int64_t offset0 = offsets.front();
+  std::for_each(offsets.begin(), offsets.end(),
+                [&offset0](std::int64_t& v) { v -= offset0; });
+  io::hdf5::close_file(h5file);
+
+  // Create cell topologies for each celltype in mesh
+  std::vector<std::vector<std::int64_t>> cells_local(recv_types.size());
+  for (std::size_t j = 0; j < types.size(); ++j)
+  {
+    std::int32_t type_index = type_to_index.at(types[j]);
+    mesh::CellType cell_type = dolfinx_cell_type[type_index];
+    auto perm = io::cells::perm_vtk(cell_type, offsets[j + 1] - offsets[j]);
+
+    for (std::size_t k = 0; k < offsets[j + 1] - offsets[j]; ++k)
+      cells_local[type_index].push_back(topology[perm[k] + offsets[j]]);
+  }
+
+  // Make first order coordinate elements
+  std::vector<fem::CoordinateElement<U>> coordinate_elements;
+  for (auto& cell : dolfinx_cell_type)
+    coordinate_elements.push_back(fem::CoordinateElement<U>(cell, 1));
+
+  auto part = create_cell_partitioner(mesh::GhostMode::none);
+  std::vector<std::span<const std::int64_t>> cells_span;
+  for (auto& cells : cells_local)
+    cells_span.push_back(cells);
+  std::array<std::size_t, 2> xs
+      = {(std::size_t)x_shape[0], (std::size_t)x_shape[1]};
+  return create_mesh(comm, comm, cells_span, coordinate_elements, comm,
+                     points_local, xs, part);
+}

cpp/dolfinx/io/cells.cpp

@@ -703,6 +703,10 @@ std::int8_t io::cells::get_vtk_cell_type(mesh::CellType cell, int dim)
     return 71;
   case mesh::CellType::hexahedron:
     return 72;
+  case mesh::CellType::pyramid:
+    return 14;
+  case mesh::CellType::prism:
+    return 73;
   default:
     throw std::runtime_error("Unknown cell type");
   }

cpp/dolfinx/io/dolfinx_io.h

@@ -11,3 +11,4 @@ namespace dolfinx::io
 
 #include <dolfinx/io/ADIOS2Writers.h>
 #include <dolfinx/io/VTKFile.h>
+#include <dolfinx/io/VTKHDF.h>

python/dolfinx/io/__init__.py

@@ -6,10 +6,10 @@
 """Tools for file input/output (IO)."""
 
 from dolfinx import cpp as _cpp
-from dolfinx.io import gmshio
+from dolfinx.io import gmshio, vtkhdf
 from dolfinx.io.utils import VTKFile, XDMFFile, distribute_entity_data
 
-__all__ = ["VTKFile", "XDMFFile", "distribute_entity_data", "gmshio"]
+__all__ = ["VTKFile", "XDMFFile", "distribute_entity_data", "gmshio", "vtkhdf"]
 
 if _cpp.common.has_adios2:
     # VTXWriter requires ADIOS2