Add function to read CG1 functions from VTKHDF file

Author: mleoni-pf

cpp/dolfinx/io/VTKHDF.h

@@ -578,4 +578,138 @@ std::vector<U> read_data(std::string point_or_cell, std::string filename,
 
   return values;
 }
+
+/// @brief Read a CG1 function from a VTKHDF format file.
+///
+/// @tparam U Scalar type of mesh
+/// @param filename Name of the file to read from.
+/// @param mesh Mesh previously read from the same file.
+/// @param timestep The time step to read for time-dependent data.
+/// @return The function read from file.
+///
+/// @note This only supports meshes with a single cell type as of now.
+template <std::floating_point U>
+fem::Function<U> read_CG1_function(std::string filename,
+                                   std::shared_ptr<mesh::Mesh<U>> mesh,
+                                   std::int32_t timestep = 0)
+{
+  auto element = basix::create_element<U>(
+      basix::element::family::P,
+      mesh::cell_type_to_basix_type(mesh->topology()->cell_type()), 1,
+      basix::element::lagrange_variant::unset,
+      basix::element::dpc_variant::unset, false);
+
+  hid_t h5file = hdf5::open_file(mesh->comm(), filename, "r", true);
+  std::string dataset_name = "/VTKHDF/PointData/u";
+
+  std::vector<std::int64_t> shape
+      = hdf5::get_dataset_shape(h5file, dataset_name);
+  hdf5::close_file(h5file);
+
+  const std::size_t block_size = shape[1];
+
+  auto P1
+      = std::make_shared<fem::FunctionSpace<U>>(fem::create_functionspace<U>(
+          mesh, std::make_shared<fem::FiniteElement<U>>(
+                    element, std::vector<std::size_t>{block_size})));
+
+  fem::Function<U> u_in(P1);
+
+  std::shared_ptr<const common::IndexMap> index_map(
+      mesh->geometry().index_map());
+
+  std::int64_t range0 = index_map->local_range()[0];
+
+  int npoints = index_map->size_local();
+
+  std::array<std::int64_t, 2> range{range0, range0 + npoints};
+
+  const auto values
+      = io::VTKHDF::read_data("Point", filename, *mesh, range, timestep);
+
+  // Parallel distribution. For vector functions we distribute each component
+  // separately.
+  std::vector<std::vector<double>> values_s(block_size);
+  for (auto& v : values_s)
+  {
+    v.reserve(values.size() / block_size);
+  }
+  for (std::size_t i = 0; i < values.size() / block_size; ++i)
+  {
+    for (int j = 0; j < block_size; ++j)
+    {
+      values_s[j].push_back(values[block_size * i + j]);
+    }
+  }
+
+  std::vector<std::int64_t> entities(range[1] - range[0]);
+  std::iota(entities.begin(), entities.end(), range[0]);
+
+  MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
+      const std::int64_t,
+      MDSPAN_IMPL_STANDARD_NAMESPACE::dextents<std::size_t, 2>>
+      entities_span(entities.data(),
+                    std::array<std::size_t, 2>{entities.size(), 1});
+
+  std::vector<std::pair<std::vector<std::int32_t>, std::vector<double>>>
+      entities_values;
+
+  for (std::size_t i = 0; i < values_s.size(); ++i)
+  {
+    entities_values.push_back(dolfinx::io::distribute_entity_data<double>(
+        *mesh->topology(), mesh->geometry().input_global_indices(),
+        mesh->geometry().index_map()->size_global(),
+        mesh->geometry().cmap().create_dof_layout(), mesh->geometry().dofmap(),
+        mesh::cell_dim(mesh::CellType::point), entities_span, values_s[i]));
+  }
+
+  auto num_vertices_per_cell
+      = dolfinx::mesh::num_cell_vertices(mesh->topology()->cell_type());
+  std::vector<std::int32_t> local_vertex_map(num_vertices_per_cell);
+
+  for (int i = 0; i < num_vertices_per_cell; ++i)
+  {
+    const auto v_to_d
+        = u_in.function_space()->dofmap()->element_dof_layout().entity_dofs(0,
+                                                                            i);
+    assert(v_to_d.size() == 1);
+    local_vertex_map[i] = v_to_d.front();
+  }
+
+  const auto tdim = mesh->topology()->dim();
+  const auto c_to_v = mesh->topology()->connectivity(tdim, 0);
+  std::vector<std::int32_t> vertex_to_dofmap(
+      mesh->topology()->index_map(0)->size_local()
+      + mesh->topology()->index_map(0)->num_ghosts());
+
+  for (int i = 0; i < mesh->topology()->index_map(tdim)->size_local(); ++i)
+  {
+    const auto local_vertices = c_to_v->links(i);
+    const auto local_dofs = u_in.function_space()->dofmap()->cell_dofs(i);
+    for (int j = 0; j < num_vertices_per_cell; ++j)
+    {
+      vertex_to_dofmap[local_vertices[j]] = local_dofs[local_vertex_map[j]];
+    }
+  }
+
+  /*
+   * After the data is read and distributed, we need to place the
+   * retrieved values in the correct position in the function's array,
+   * reading values and positions from `entities_values`.
+   */
+  for (std::size_t i = 0; i < entities_values[0].first.size(); ++i)
+  {
+    for (std::size_t j = 0; j < block_size; ++j)
+    {
+      u_in.x()
+          ->array()[block_size * vertex_to_dofmap[entities_values[0].first[i]]
+                    + j]
+          = entities_values[j].second[i];
+    }
+  }
+
+  u_in.x()->scatter_fwd();
+
+  return u_in;
+}
 } // namespace dolfinx::io::VTKHDF

python/dolfinx/io/vtkhdf.py

@@ -12,19 +12,22 @@
 import numpy.typing as npt
 
 import basix
+import dolfinx
 import ufl
 from dolfinx.cpp.io import (
+    read_vtkhdf_cg1_function,
     read_vtkhdf_mesh_float32,
     read_vtkhdf_mesh_float64,
     write_vtkhdf_cg1_function,
     write_vtkhdf_data,
     write_vtkhdf_dg0_function,
     write_vtkhdf_mesh,
 )
-from dolfinx.fem import Function
+from dolfinx.fem import Function, FunctionSpace
 from dolfinx.mesh import Mesh
 
 __all__ = [
+    "read_cg1_function",
     "read_mesh",
     "write_cell_data",
     "write_cg1_function",
@@ -134,3 +137,33 @@ def write_dg0_function(filename: str | Path, mesh: Mesh, u: Function, time: floa
         time: Timestamp.
     """
     write_vtkhdf_dg0_function(filename, mesh._cpp_object, u._cpp_object, time)
+
+
+def read_cg1_function(filename: str | Path, mesh: Mesh, timestep: int = 0) -> Function:
+    """Read a CG1 function form file.
+
+    Args:
+        filename: File to read.
+        mesh: Mesh.
+        timestep: Time step to read
+
+    Returns:
+        The read function.
+    """
+    match type(mesh._cpp_object):
+        case dolfinx.cpp.mesh.Mesh_float64:
+            dtype = np.float64
+        case dolfinx.cpp.mesh.Mesh_float32:
+            dtype = np.float32
+    u_cpp = read_vtkhdf_cg1_function(filename, mesh._cpp_object, timestep)
+    elem = basix.ufl.element(
+        "Lagrange",
+        mesh.basix_cell(),
+        1,
+        shape=tuple(u_cpp.function_space.element.value_shape),
+        dtype=dtype,
+    )
+    V = FunctionSpace(mesh, elem, u_cpp.function_space)
+    u_ret = Function(V, dtype=dtype)
+    u_ret.interpolate(u_cpp)
+    return u_ret

python/dolfinx/wrappers/io.cpp

@@ -229,6 +229,8 @@ void io(nb::module_& m)
   m.def("write_vtkhdf_cg1_function", &dolfinx::io::VTKHDF::write_CG1_function<double>);
   m.def("write_vtkhdf_dg0_function", &dolfinx::io::VTKHDF::write_DG0_function<float>);
   m.def("write_vtkhdf_dg0_function", &dolfinx::io::VTKHDF::write_DG0_function<double>);
+  m.def("read_vtkhdf_cg1_function", &dolfinx::io::VTKHDF::read_CG1_function<float>);
+  m.def("read_vtkhdf_cg1_function", &dolfinx::io::VTKHDF::read_CG1_function<double>);
   m.def("read_vtkhdf_mesh_float64",
         [](MPICommWrapper comm, const std::string& filename, std::size_t gdim,
            std::optional<std::int32_t> max_facet_to_cell_links)

python/test/unit/io/test_vtkhdf.py

@@ -9,9 +9,17 @@
 import numpy as np
 import pytest
 
+import basix
 import dolfinx
 import ufl
-from dolfinx.io.vtkhdf import read_mesh, write_cell_data, write_mesh, write_point_data
+from dolfinx.io.vtkhdf import (
+    read_cg1_function,
+    read_mesh,
+    write_cell_data,
+    write_cg1_function,
+    write_mesh,
+    write_point_data,
+)
 from dolfinx.mesh import CellType, Mesh, create_unit_cube, create_unit_square
 
 
@@ -206,3 +214,43 @@ def test_write_mixed_topology_data(mixed_topology_mesh):
     b = sum([im.size_local for im in mesh.topology.index_maps(mesh.topology.dim)])
     cell_data = np.arange(b, dtype=np.float64)
     write_cell_data(filename, mesh, cell_data, 0.0)
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("width", [1, 3])
+def test_write_read_cg1_function(dtype, width):
+    mesh = create_unit_cube(MPI.COMM_WORLD, 5, 5, 5, dtype=dtype)
+    elem = basix.ufl.element("Lagrange", mesh.basix_cell(), 1, shape=(width,), dtype=dtype)
+    V = dolfinx.fem.functionspace(mesh, elem)
+    f = dolfinx.fem.Function(V, dtype=dtype)
+    f.interpolate(lambda x: np.vstack([np.sin(x[0]) ** 2 + np.cos(x[1] ** 2 * x[2])] * width))
+
+    filename = f"cg1_fun_{width}.vtkhdf"
+    write_mesh(filename, mesh)
+    write_cg1_function(filename, mesh, f, 0)
+
+    mesh2 = read_mesh(mesh.comm, filename, dtype=dtype)
+    f_2 = read_cg1_function(filename, mesh2)
+
+    filename2 = f"cg1_fun2_{width}.vtkhdf"
+    write_mesh(filename2, mesh2)
+    write_cg1_function(filename2, mesh2, f_2, 0)
+
+    cell_map = mesh.topology.index_map(mesh.topology.dim)
+    num_cells_on_proc = cell_map.size_local + cell_map.num_ghosts
+    cells = np.arange(num_cells_on_proc, dtype=np.int32)
+    padding = 1e-14 if dtype == np.float64 else 1e-5
+    interpolation_data = dolfinx.fem.create_interpolation_data(
+        V, f_2.function_space, cells, padding=padding
+    )
+
+    f_i = dolfinx.fem.Function(V, dtype=dtype)
+    f_i.interpolate_nonmatching(f_2, cells, interpolation_data=interpolation_data)
+    f_i.x.scatter_forward()
+
+    filename_i = f"cg1_funi_{width}.vtkhdf"
+    write_mesh(filename_i, mesh)
+    write_cg1_function(filename_i, mesh, f_i, 0)
+
+    tol = 1e-14 if dtype == np.float64 else 1e-6
+    np.testing.assert_allclose(f.x.array, f_i.x.array, rtol=tol)