Resolve jorgensd/adios4dolfinx#17 by adding interfaces for legacy meshtags in both adios2 and h5py backend

Author: jorgensd

src/adios4dolfinx/backends/adios2/backend.py

@@ -31,6 +31,8 @@ def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]
     args = arguments or {}
     if "engine" not in args.keys():
         args["engine"] = "BP4"
+    if "legacy" not in args.keys():
+        args["legacy"] = False  # Only used for legacy HDF5 meshtags
     return args
 
 
@@ -71,7 +73,7 @@ def write_attributes(
         filename=filename,
         mode=adios2.Mode.Append,
         io_name="AttributeWriter",
-        **backend_args,
+        engine=backend_args["engine"],
     ) as adios_file:
         adios_file.file.BeginStep()
 
@@ -105,7 +107,7 @@ def read_attributes(
         adios=adios,
         filename=filename,
         mode=adios2.Mode.Read,
-        **backend_args,
+        engine=backend_args["engine"],
         io_name="AttributesReader",
     ) as adios_file:
         adios_file.file.BeginStep()
@@ -143,7 +145,7 @@ def read_timestamps(
         adios=adios,
         filename=filename,
         mode=adios2.Mode.Read,
-        **backend_args,
+        engine=backend_args["engine"],
         io_name="TimestepReader",
     ) as adios_file:
         time_name = f"{function_name}_time"
@@ -195,7 +197,8 @@ def write_mesh(
         filename=filename,
         mode=mode,
         comm=comm,
-        **backend_args,
+        engine=backend_args["engine"],
+        io_name=backend_args["io_name"],
     ) as adios_file:
         adios_file.file.BeginStep()
         # Write geometry
@@ -473,73 +476,126 @@ def read_meshtags_data(
     """
 
     adios = adios2.ADIOS(comm)
-    backend_args = {} if backend_args is None else backend_args
+    backend_args = get_default_backend_args(backend_args)
     io_name = backend_args.get("io_name", "MeshTagsReader")
-    engine = backend_args.get("engine", "BP4")
+    engine = backend_args["engine"]
+    legacy = backend_args["legacy"]
     with ADIOSFile(
         adios=adios,
         filename=filename,
         mode=adios2.Mode.Read,
         engine=engine,
         io_name=io_name,
     ) as adios_file:
-        # Get mesh cell type
-        dim_attr_name = f"{name}_dim"
-        step = 0
-        for i in range(adios_file.file.Steps()):
-            adios_file.file.BeginStep()
-            if dim_attr_name in adios_file.io.AvailableAttributes().keys():
-                step = i
-                break
-            adios_file.file.EndStep()
-        if dim_attr_name not in adios_file.io.AvailableAttributes().keys():
-            raise KeyError(f"{dim_attr_name} not found in {filename}")
+        if not legacy:
+            # Get mesh cell type
+            dim_attr_name = f"{name}_dim"
+            step = 0
+            for i in range(adios_file.file.Steps()):
+                adios_file.file.BeginStep()
+                if dim_attr_name in adios_file.io.AvailableAttributes().keys():
+                    step = i
+                    break
+                adios_file.file.EndStep()
+            if dim_attr_name not in adios_file.io.AvailableAttributes().keys():
+                raise KeyError(f"{dim_attr_name} not found in {filename}")
 
-        m_dim = adios_file.io.InquireAttribute(dim_attr_name)
-        dim = int(m_dim.Data()[0])
+            m_dim = adios_file.io.InquireAttribute(dim_attr_name)
+            dim = int(m_dim.Data()[0])
 
-        # Get mesh tags entites
-        topology_name = f"{name}_topology"
-        for i in range(step, adios_file.file.Steps()):
-            if i > step:
-                adios_file.file.BeginStep()
-            if topology_name in adios_file.io.AvailableVariables().keys():
-                break
-            adios_file.file.EndStep()
-        if topology_name not in adios_file.io.AvailableVariables().keys():
-            raise KeyError(f"{topology_name} not found in {filename}")
+            # Get mesh tags entites
+            topology_name = f"{name}_topology"
+            for i in range(step, adios_file.file.Steps()):
+                if i > step:
+                    adios_file.file.BeginStep()
+                if topology_name in adios_file.io.AvailableVariables().keys():
+                    break
+                adios_file.file.EndStep()
+            if topology_name not in adios_file.io.AvailableVariables().keys():
+                raise KeyError(f"{topology_name} not found in {filename}")
+
+            topology = adios_file.io.InquireVariable(topology_name)
+            top_shape = topology.Shape()
+            topology_range = compute_local_range(comm, top_shape[0])
+
+            topology.SetSelection(
+                [
+                    [topology_range[0], 0],
+                    [topology_range[1] - topology_range[0], top_shape[1]],
+                ]
+            )
+            mesh_entities = np.empty(
+                (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
+            )
+            adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
+
+            # Get mesh tags values
+            values_name = f"{name}_values"
+            if values_name not in adios_file.io.AvailableVariables().keys():
+                raise KeyError(f"{values_name} not found")
+
+            values = adios_file.io.InquireVariable(values_name)
+            val_shape = values.Shape()
+            assert val_shape[0] == top_shape[0]
+            values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
+            tag_values = np.empty(
+                (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
+            )
+            adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
 
-        topology = adios_file.io.InquireVariable(topology_name)
-        top_shape = topology.Shape()
-        topology_range = compute_local_range(comm, top_shape[0])
+            adios_file.file.PerformGets()
+            adios_file.file.EndStep()
+        else:
+            # Get mesh cell type
+            dim_attr_name = f"{name}_dim"
+            assert adios_file.file.Steps() == 0
+            if (ct_key := f"/{name}/topology/celltype") in adios_file.io.AvailableAttributes():
+                cell_type = adios_file.io.InquireAttribute(ct_key)
+            else:
+                raise ValueError(f"Celltype not found for meshtags {name} in {filename}.")
+            dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type.DataString()[0]))
+
+            # Get mesh tags entites
+            if (top_key := f"/{name}/topology") in adios_file.io.AvailableVariables():
+                topology = adios_file.io.InquireVariable(top_key)
+            else:
+                raise ValueError(f"Topology not found for meshtags {name} in {filename}.")
+
+            top_shape = topology.Shape()
+            topology_range = compute_local_range(comm, top_shape[0])
+
+            topology.SetSelection(
+                [
+                    [topology_range[0], 0],
+                    [topology_range[1] - topology_range[0], top_shape[1]],
+                ]
+            )
+            mesh_entities = np.empty(
+                (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
+            )
+            adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
 
-        topology.SetSelection(
-            [
-                [topology_range[0], 0],
-                [topology_range[1] - topology_range[0], top_shape[1]],
-            ]
-        )
-        mesh_entities = np.empty(
-            (topology_range[1] - topology_range[0], top_shape[1]), dtype=np.int64
-        )
-        adios_file.file.Get(topology, mesh_entities, adios2.Mode.Deferred)
+            # Get mesh tags values
+            if (val_key := f"/{name}/values") in adios_file.io.AvailableVariables():
+                values = adios_file.io.InquireVariable(val_key)
+            else:
+                raise ValueError(f"Values not found for meshtags {name} in {filename}.")
 
-        # Get mesh tags values
-        values_name = f"{name}_values"
-        if values_name not in adios_file.io.AvailableVariables().keys():
-            raise KeyError(f"{values_name} not found")
+            val_shape = values.Shape()
+            assert val_shape[0] == top_shape[0]
 
-        values = adios_file.io.InquireVariable(values_name)
-        val_shape = values.Shape()
-        assert val_shape[0] == top_shape[0]
-        values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
-        tag_values = np.empty((topology_range[1] - topology_range[0]), dtype=np.int32)
-        adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
+            values.SetSelection([[topology_range[0]], [topology_range[1] - topology_range[0]]])
+            tag_values = np.empty(
+                (topology_range[1] - topology_range[0]), dtype=values.Type().strip("_t")
+            )
+            adios_file.file.Get(values, tag_values, adios2.Mode.Deferred)
 
-        adios_file.file.PerformGets()
-        adios_file.file.EndStep()
+            adios_file.file.PerformGets()
+            adios_file.file.EndStep()
 
-        return MeshTagsData(name=name, values=tag_values, indices=mesh_entities, dim=dim)
+        return MeshTagsData(
+            name=name, values=tag_values.astype(np.int32), indices=mesh_entities, dim=dim
+        )
 
 
 def read_dofmap(

src/adios4dolfinx/backends/h5py/backend.py

@@ -13,6 +13,7 @@
 from mpi4py import MPI
 
 import dolfinx
+import h5py
 import numpy as np
 import numpy.typing as npt
 from dolfinx.graph import adjacencylist
@@ -23,6 +24,10 @@
 
 read_mode = ReadMode.parallel
 
+# try:
+# except ModuleNotFoundError:
+#     raise ModuleNotFoundError("This backend requires h5py to be installed.")
+
 
 @contextlib.contextmanager
 def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None):
@@ -35,8 +40,6 @@ def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None):
         comm: The MPI communicator
 
     """
-    import h5py
-
     if comm is None:
         comm = MPI.COMM_WORLD
 
@@ -58,13 +61,7 @@ def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None):
 
 
 def get_default_backend_args(arguments: dict[str, Any] | None) -> dict[str, Any]:
-    args = arguments or {}
-
-    if arguments:
-        # Currently no default arguments for h5py backend
-        # TODO: Pehaps we would like to make this into a warning instead?
-        raise RuntimeError("Unexpected backend arguments to h5py backend")
-
+    args = arguments or {"legacy": False}  # If meshtags is read from legacy
     return args
 
 
@@ -384,20 +381,32 @@ def read_meshtags_data(
         Internal data structure for the mesh tags read from file
     """
     backend_args = get_default_backend_args(backend_args)
+    legacy = backend_args["legacy"]
     with h5pyfile(filename, filemode="r", comm=comm, force_serial=False) as h5file:
-        if "mesh" not in h5file.keys():
-            raise KeyError("No mesh found")
-        mesh = h5file["mesh"]
-        if "tags" not in mesh.keys():
-            raise KeyError("Could not find 'tags' in file, are you sure this is a checkpoint?")
-        tags = mesh["tags"]
-        if name not in tags.keys():
-            raise KeyError(f"Could not find {name} in '/mesh/tags/' in {filename}")
-        tag = tags[name]
-
-        dim = tag.attrs["dim"]
-        topology = tag["Topology"]
-        values = tag["Values"]
+        if legacy:
+            if name not in h5file.keys():
+                raise RuntimeError(f"MeshTag {name} not found in {filename}.")
+            mesh = h5file[name]
+            topology = mesh["topology"]
+            cell_type = topology.attrs["celltype"]
+            if isinstance(cell_type, np.bytes_):
+                cell_type = cell_type.decode("utf-8")
+            dim = dolfinx.mesh.cell_dim(dolfinx.mesh.to_type(cell_type))
+            values = mesh["values"]
+        else:
+            if "mesh" not in h5file.keys():
+                raise KeyError("No mesh found")
+            mesh = h5file["mesh"]
+            if "tags" not in mesh.keys():
+                raise KeyError("Could not find 'tags' in file, are you sure this is a checkpoint?")
+            tags = mesh["tags"]
+            if name not in tags.keys():
+                raise KeyError(f"Could not find {name} in '/mesh/tags/' in {filename}")
+            tag = tags[name]
+
+            dim = tag.attrs["dim"]
+            topology = tag["Topology"]
+            values = tag["Values"]
         num_entities_global = topology.shape[0]
         topology_range = compute_local_range(comm, num_entities_global)
         indices = topology[slice(*topology_range), :].astype(np.int64)

src/adios4dolfinx/backends/pyvista/backend.py

@@ -10,7 +10,7 @@
 try:
     import pyvista
 except ImportError:
-    raise ModuleNotFoundError("This module requires pyvista")
+    raise ModuleNotFoundError("The PyVista-backend requires pyvista")
 from pathlib import Path
 
 from mpi4py import MPI
@@ -217,10 +217,10 @@ def read_point_data(
         else:
             num_components = dataset.shape[1]
         if np.issubdtype(dataset.dtype, np.integer):
-            gtype = in_data.points.dtype
+            gtype = grid.points.dtype
             dataset = dataset.astype(gtype)
         else:
-            gtype = in_data.dtype
+            gtype = dataset.dtype
         num_components, gtype = comm.bcast((num_components, gtype), root=0)
         local_range_start = 0
     else:

src/adios4dolfinx/backends/vtkhdf/backend.py

@@ -17,7 +17,7 @@
 from adios4dolfinx.utils import check_file_exists, compute_local_range
 
 from .. import FileMode, ReadMode
-from ..h5py.backend import convert_file_mode, h5pyfile
+from .._h5py.backend import convert_file_mode, h5pyfile
 from ..pyvista.backend import _arbitrary_lagrange_vtk, _cell_degree, _first_order_vtk
 
 read_mode = ReadMode.parallel

src/adios4dolfinx/backends/xdmf/backend.py

@@ -2,7 +2,6 @@
 Module that uses DOLFINx/H%py to import XDMF files.
 """
 
-import contextlib
 from pathlib import Path
 from typing import Any
 from xml.etree import ElementTree
@@ -18,43 +17,11 @@
 from adios4dolfinx.utils import check_file_exists, compute_local_range
 
 from .. import FileMode, ReadMode
+from ..h5py.backend import h5pyfile
 
 read_mode = ReadMode.parallel
 
 
-@contextlib.contextmanager
-def h5pyfile(h5name, filemode="r", force_serial: bool = False, comm=None):
-    """Context manager for opening an HDF5 file with h5py.
-
-    Args:
-        h5name: The name of the HDF5 file.
-        filemode: The file mode.
-        force_serial: Force serial access to the file.
-        comm: The MPI communicator
-
-    """
-    import h5py
-
-    if comm is None:
-        comm = MPI.COMM_WORLD
-
-    if h5py.h5.get_config().mpi and not force_serial:
-        h5file = h5py.File(h5name, filemode, driver="mpio", comm=comm)
-    else:
-        if comm.size > 1 and not force_serial:
-            raise ValueError(
-                f"h5py is not installed with MPI support, while using {comm.size} processes.",
-                "If you really want to do this, turn on the `force_serial` flag.",
-            )
-        h5file = h5py.File(h5name, filemode)
-
-    try:
-        yield h5file
-    finally:
-        if h5file.id:
-            h5file.close()
-
-
 def extract_function_names_and_timesteps(filename: Path | str) -> dict[str, list[str]]:
     tree = ElementTree.parse(filename)
     root = tree.getroot()