Add vtkhdf reader

Author: jorgensd

src/adios4dolfinx/backends/__init__.py

@@ -151,7 +151,7 @@ def read_mesh_data(
         self,
         filename: Path | str,
         comm: MPI.Intracomm,
-        time: float | None,
+        time: str | float | None,
         read_from_partition: bool,
         backend_args: dict[str, Any] | None,
     ) -> ReadMeshData:

src/adios4dolfinx/backends/adios2/backend.py

@@ -275,7 +275,7 @@ def write_mesh(
 def read_mesh_data(
     filename: Path | str,
     comm: MPI.Intracomm,
-    time: float | None = 0.0,
+    time: str | float | None = 0.0,
     read_from_partition: bool = False,
     backend_args: dict[str, Any] | None = None,
 ) -> ReadMeshData:

src/adios4dolfinx/backends/h5py/backend.py

@@ -239,7 +239,7 @@ def write_mesh(
 def read_mesh_data(
     filename: Path | str,
     comm: MPI.Intracomm,
-    time: float | None = 0.0,
+    time: str | float | None = 0.0,
     read_from_partition: bool = False,
     backend_args: dict[str, Any] | None = None,
 ) -> ReadMeshData:

src/adios4dolfinx/backends/pyvista/backend.py

@@ -109,7 +109,7 @@ def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]
 def read_mesh_data(
     filename: Path | str,
     comm: MPI.Intracomm,
-    time: float | None,
+    time: str | float | None,
     read_from_partition: bool,
     backend_args: dict[str, Any] | None,
 ) -> ReadMeshData:

src/adios4dolfinx/backends/vtkhdf/backend.py

@@ -89,20 +89,71 @@ def read_mesh_data(
 
     with h5pyfile(filename, "r", comm=comm) as h5file:
         hdf = h5file["VTKHDF"]
-        num_cells_global = hdf["Types"].size
 
-        local_cell_range = compute_local_range(comm, num_cells_global)
-        cell_types_local = hdf["Types"][local_cell_range[0] : local_cell_range[1]]
+        if time is None:
+            num_cells_global = hdf["Types"].size
+            local_cell_range = compute_local_range(comm, num_cells_global)
+            cell_types_local = hdf["Types"][slice(*local_cell_range)]
+
+            num_points_global = hdf["NumberOfPoints"][0]
+            local_point_range = compute_local_range(comm, num_points_global)
+            points_local = hdf["Points"][slice(*local_point_range)]
+
+            # Connectivity read
+            offsets = hdf["Offsets"]
+            local_connectivity_offset = offsets[local_cell_range[0] : local_cell_range[1] + 1]
+            topology = hdf["Connectivity"][
+                local_connectivity_offset[0] : local_connectivity_offset[-1]
+            ]
+            offset = local_connectivity_offset - local_connectivity_offset[0]
+        else:
+            if "Steps" not in hdf.keys():
+                raise RuntimeError(f"No timestepping information found in {filename}.")
+            stamps = hdf["Steps"]["Values"][:]
+            pos = np.flatnonzero(np.isclose(stamps, time))
+            if len(pos) == 0:
+                raise RuntimeError(f"Could not find mesh at t={time} in {filename}.")
+            elif len(pos) > 1:
+                raise RuntimeError(f"Multiple time steps for mesh at t={time} in {filename}")
 
-        num_points_global = hdf["NumberOfPoints"][0]
-        local_point_range = compute_local_range(comm, num_points_global)
-        points_local = hdf["Points"][local_point_range[0] : local_point_range[1]]
+            # Get number of points
+            point_node = hdf["Points"]
+            step_start = hdf["Steps"]["PointOffsets"][pos[0]]
 
-        # Connectivity read
-        offsets = hdf["Offsets"]
-        local_connectivity_offset = offsets[local_cell_range[0] : local_cell_range[1] + 1]
-        topology = hdf["Connectivity"][local_connectivity_offset[0] : local_connectivity_offset[-1]]
-    offset = local_connectivity_offset - local_connectivity_offset[0]
+            # NOTE: currently, it doesn't seem like we follow:
+            # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata
+            # As only one num_points is stored irregardless of time steps added.
+            if hdf["NumberOfPoints"].shape[0] != len(stamps):
+                num_pts = hdf["NumberOfPoints"][0]
+            else:
+                num_pts = hdf["NumberOfPoints"][pos[0]]
+            lr = compute_local_range(comm, num_pts)
+            points_local = point_node[step_start + lr[0] : step_start + lr[1]]
+
+            # Get cell-types in step
+            cell_start = hdf["Steps"]["CellOffsets"][pos[0]]
+            if hdf["NumberOfCells"].shape[0] != len(stamps):
+                num_cells = hdf["NumberOfCells"][0]
+            else:
+                num_cells = hdf["NumberOfCells"][pos[0]]
+            local_cell_range = compute_local_range(comm, num_cells)
+            cell_types_local = hdf["Types"][
+                cell_start + local_cell_range[0] : cell_start + local_cell_range[1]
+            ]
+
+            # Get connectivity in step
+            connectivity_start = hdf["Steps"]["ConnectivityIdOffsets"][pos[0]]
+            # Connectivity read
+            offsets = hdf["Offsets"]
+            local_connectivity_offset = offsets[
+                connectivity_start + local_cell_range[0] : connectivity_start
+                + local_cell_range[1]
+                + 1
+            ]
+            topology = hdf["Connectivity"][
+                local_connectivity_offset[0] : local_connectivity_offset[-1]
+            ]
+            offset = local_connectivity_offset - local_connectivity_offset[0]
 
     # NOTE: Currently we limit ourselfs to a single celltype, as it makes life easier,
     # other things have to change in `MeshReadData` to support this.
@@ -172,13 +223,15 @@ def read_point_data(
                 raise RuntimeError(f"Could not find {name}(t={time}) in {filename}.")
             elif len(pos) > 1:
                 raise RuntimeError(f"Multiple time steps for {name}(t={time}) in {filename}")
-            pd_offsets = hdf["Steps"]["PointDataOffsets"][name][:]
-            pd_offsets = np.hstack([pd_offsets, [func_node.shape[0]]])
-
-            step_start = pd_offsets[pos[0]]
-            step_end = pd_offsets[pos[0] + 1]
-            global_size = step_end - step_start
-            lr = compute_local_range(comm, global_size)
+            step_start = hdf["Steps"]["PointDataOffsets"][name][pos[0]]
+            # NOTE: currently, it doesn't seem like we follow:
+            # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata
+            # As only one num_points is stored irregardless of time steps added.
+            if hdf["NumberOfPoints"].shape[0] != len(stamps):
+                num_pts = hdf["NumberOfPoints"][0]
+            else:
+                num_pts = hdf["NumberOfPoints"][pos[0]]
+            lr = compute_local_range(comm, num_pts)
             return func_node[step_start + lr[0] : step_start + lr[1]], lr[0]
 
 
@@ -197,12 +250,37 @@ def read_cell_data(
             raise RuntimeError(f"No cell data found in {filename}.")
         cell_data = hdf["CellData"]
         assert cell_data is not None
+        if name not in cell_data.keys():
+            raise ValueError(f"No cell data with name {name} in {filename}")
         cell_data_node = cell_data[name]
-        cell_data_shape = cell_data_node.shape
-        num_cells_global = hdf["Types"].size
-        assert num_cells_global == cell_data_shape[0]
-        local_cell_range = compute_local_range(comm, num_cells_global)
-        data = cell_data_node[slice(*local_cell_range)]
+        assert cell_data_node is not None
+        if time is None:
+            cell_data_shape = cell_data_node.shape
+            num_cells_global = hdf["Types"].size
+            assert num_cells_global == cell_data_shape[0]
+            local_cell_range = compute_local_range(comm, num_cells_global)
+            data = cell_data_node[slice(*local_cell_range)]
+        else:
+            if "Steps" not in hdf.keys():
+                raise RuntimeError(f"No timestepping information found in {filename}.")
+            stamps = hdf["Steps"]["Values"][:]
+            pos = np.flatnonzero(np.isclose(stamps, time))
+            if len(pos) == 0:
+                raise RuntimeError(f"Could not find {name}(t={time}) in {filename}.")
+            elif len(pos) > 1:
+                raise RuntimeError(f"Multiple time steps for {name}(t={time}) in {filename}")
+            cd_start = hdf["Steps"]["CellDataOffsets"][name][pos[0]]
+
+            # NOTE: currently, it doesn't seem like we follow:
+            # https://docs.vtk.org/en/latest/vtk_file_formats/vtkhdf_file_format/vtkhdf_specifications.html#temporal-unstructuredgrid-and-polydata
+            # As only one num_points is stored irregardless of time steps added.
+            if hdf["NumberOfCells"].shape[0] != len(stamps):
+                number_of_cells = hdf["NumberOfCells"][0]
+            else:
+                number_of_cells = hdf["NumberOfCells"][pos[0]]
+            lr = compute_local_range(comm, number_of_cells)
+            data = cell_data_node[cd_start + lr[0] : cd_start + lr[1]]
+
     # NOTE: THis could be optimized by hand-coding some communication in
     # `read_cell_data` on the frontend side
     md = read_mesh_data(filename, comm, time=time, read_from_partition=False, backend_args=None)

src/adios4dolfinx/backends/xdmf/backend.py

@@ -93,7 +93,7 @@ def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]
 def read_mesh_data(
     filename: Path | str,
     comm: MPI.Intracomm,
-    time: float | None,
+    time: str | float | None,
     read_from_partition: bool,
     backend_args: dict[str, Any] | None,
 ) -> ReadMeshData:
@@ -109,6 +109,7 @@ def read_mesh_data(
     Returns:
         Internal data structure for the mesh data read from file
     """
+    assert read_from_partition == False
     check_file_exists(filename)
     with dolfinx.io.XDMFFile(comm, filename, "r") as file:
         cell_shape, cell_degree = file.read_cell_type()

src/adios4dolfinx/checkpointing.py

@@ -378,7 +378,7 @@ def read_mesh(
     filename: Path | str,
     comm: MPI.Intracomm,
     ghost_mode: dolfinx.mesh.GhostMode = dolfinx.mesh.GhostMode.shared_facet,
-    time: float = 0.0,
+    time: float | str | None = None,
     read_from_partition: bool = False,
     backend_args: dict[str, Any] | None = None,
     backend: str = "adios2",

tests/test_vtkhdf.py

@@ -3,8 +3,8 @@
 import numpy as np
 import pytest
 from dolfinx.fem import Function
-from dolfinx.io.vtkhdf import write_mesh, write_point_data
-from dolfinx.mesh import create_unit_cube
+from dolfinx.io.vtkhdf import write_cell_data, write_mesh, write_point_data
+from dolfinx.mesh import CellType, compute_midpoints, create_unit_cube
 
 import adios4dolfinx
 
@@ -17,18 +17,19 @@ def g(x, t):
     return x[0], 2 * x[1], -x[2] * t
 
 
+@pytest.mark.parametrize("cell_type", [CellType.hexahedron, CellType.tetrahedron])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_write_point_data(dtype, tmp_path):
+def test_write_point_data(dtype, tmp_path, cell_type):
     comm = MPI.COMM_WORLD
     tmp_path = comm.bcast(tmp_path, root=0)
     comm.barrier()
 
     # Write temporal data
-    mesh = create_unit_cube(comm, 5, 5, 5, dtype=dtype)
-    filename = tmp_path / "point_data.vtkhdf"
+    mesh = create_unit_cube(comm, 5, 5, 5, dtype=dtype, cell_type=cell_type)
+    filename = tmp_path / f"point_data_{cell_type.name}.vtkhdf"
     write_mesh(str(filename), mesh)
 
-    t = np.linspace(0, 1.2, 25)
+    t = np.linspace(0.1, 1.2, 25)
     num_nodes_local = mesh.geometry.index_map().size_local
     for ti in t:
         point_data = f(mesh.geometry.x.T[:, :num_nodes_local], ti)
@@ -44,7 +45,7 @@ def test_write_point_data(dtype, tmp_path):
             filename=filename, name="u", mesh=grid, time=tj, backend="vtkhdf"
         )
         v_ref = Function(u.function_space, dtype=u.x.array.dtype)
-        atol = 5 * np.finfo(u.x.array.dtype).eps
+        atol = 10 * np.finfo(u.x.array.dtype).eps
         v_ref.interpolate(lambda x: f(x, tj))
         np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol)
 
@@ -63,20 +64,61 @@ def test_write_point_data(dtype, tmp_path):
             filename=blocked_file, name="u", mesh=grid, time=tk, backend="vtkhdf"
         )
         v_ref = Function(u.function_space, dtype=u.x.array.dtype)
-        atol = 5 * np.finfo(u.x.array.dtype).eps
+        atol = 10 * np.finfo(u.x.array.dtype).eps
         v_ref.interpolate(lambda x: g(x, tk))
         np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol)
 
 
-# @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-# def test_write_cell_data(dtype, tmp_path):
-#     comm = MPI.COMM_WORLD
-#     tmp_path = comm.bcast(tmp_path, root=0)
-#     comm.barrier()
+@pytest.mark.parametrize("cell_type", [CellType.hexahedron, CellType.tetrahedron])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_write_cell_data(dtype, tmp_path, cell_type):
+    comm = MPI.COMM_WORLD
+    tmp_path = comm.bcast(tmp_path, root=0)
+    comm.barrier()
+
+    # Write temporal data
+    mesh = create_unit_cube(comm, 5, 5, 5, dtype=dtype, cell_type=cell_type)
+    filename = tmp_path / f"cell_data_{cell_type.name}.vtkhdf"
+    write_mesh(str(filename), mesh)
+
+    t = np.linspace(0.1, 1.2, 25)
+    num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
+    midpoints = compute_midpoints(
+        mesh, mesh.topology.dim, np.arange(num_cells_local, dtype=np.int32)
+    )
+    for ti in t:
+        cell_data = f(midpoints.T, ti)
+        write_cell_data(str(filename), mesh, cell_data, float(ti))
+    comm.barrier()
 
-#     mesh = create_unit_cube(MPI.COMM_WORLD, 5, 5, 5, dtype=dtype, cell_type=CellType.hexahedron)
-#     filename = "cell_data.vtkhdf"
-#     write_mesh(filename, mesh)
-#     cell_data = np.arange(mesh.topology.index_map(2).size_local * width)
-#     for j in range(3):
-#         write_cell_data(filename, mesh, cell_data, float(j))
+    grid = adios4dolfinx.read_mesh(filename=filename, comm=comm, backend="vtkhdf")
+    # Since we shuffle time we need to shuffle in the same way on each process
+    np.random.shuffle(t)
+    t = comm.bcast(t, root=0)
+    for tj in t:
+        u = adios4dolfinx.read_cell_data(
+            filename=filename, name="u", mesh=grid, time=tj, backend="vtkhdf"
+        )
+        v_ref = Function(u.function_space, dtype=u.x.array.dtype)
+        atol = 10 * np.finfo(u.x.array.dtype).eps
+        v_ref.interpolate(lambda x: f(x, tj))
+        np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol)
+
+    # Test blocked data as well (with shuffled input timestep)
+    blocked_file = filename.with_stem(filename.stem + "_blocked")
+    write_mesh(str(blocked_file), mesh)
+    for tj in t:
+        cell_data = np.asarray(g(midpoints.T, tj)).T.flatten()
+        write_cell_data(str(blocked_file), mesh, cell_data, float(tj))
+    comm.barrier()
+
+    np.random.shuffle(t)
+    t = comm.bcast(t, root=0)
+    for tk in t:
+        u = adios4dolfinx.read_cell_data(
+            filename=blocked_file, name="u", mesh=grid, time=tk, backend="vtkhdf"
+        )
+        v_ref = Function(u.function_space, dtype=u.x.array.dtype)
+        atol = 10 * np.finfo(u.x.array.dtype).eps
+        v_ref.interpolate(lambda x: g(x, tk))
+        np.testing.assert_allclose(u.x.array, v_ref.x.array, atol=atol)