I/O from XDMF/HDF5 files in dolfin-x

@rosilho I had some spare time today and made a “quick and dirty” solution to this relying on my adios4dolfinx library (GitHub - jorgensd/adios4dolfinx: Interface of ADIOS2 for DOLFINx) and h5py.

I tested it with the file you supplied (in serial and parallel), and got the expected results:

from pathlib import Path

import adios4dolfinx
import dolfinx
import numpy as np
from mpi4py import MPI

with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "idealized_LV_ref_0.xdmf", "r") as xdmf:
    mesh = xdmf.read_mesh(name="mesh", xpath="Xdmf/Domain/Grid")

import basix
import h5py


def read_mesh_data(file_path:Path, mesh: dolfinx.mesh.Mesh, data_path:str):
    assert file_path.is_file(), f"File {file_path} does not exist"
    infile = h5py.File(file_path, "r", driver="mpio", comm=mesh.comm)
    num_nodes_global = mesh.geometry.index_map().size_global
    assert data_path in infile.keys(), f"Data {data_path} does not exist"
    dataset = infile[data_path]
    shape = dataset.shape
    assert shape[0] == num_nodes_global, f"Got data of shape {shape}, expected {num_nodes_global, shape[1]}"
    dtype = dataset.dtype
    # Read data locally on each process
    local_input_range = adios4dolfinx.comm_helpers.compute_local_range(mesh.comm, num_nodes_global)    
    local_input_data = dataset[local_input_range[0]:local_input_range[1]]

    # Create appropriate function space (based on coordinate map)
    assert len(mesh.geometry.cmaps) == 1, "Mixed cell-type meshes not supported"
    element = basix.ufl.element(
        basix.ElementFamily.P,
        mesh.topology.cell_name(),
        mesh.geometry.cmaps[0].degree,
        mesh.geometry.cmaps[0].variant,
        shape=(shape[1],),
        gdim=mesh.geometry.dim)

    # Assumption: Same doflayout for geometry and function space, cannot test in python    
    V = dolfinx.fem.FunctionSpace(mesh, element)
    uh = dolfinx.fem.Function(V, name=data_path)
    # Assume that mesh is first order for now
    assert mesh.geometry.cmaps[0].degree == 1, "Only linear meshes supported"
    x_dofmap = mesh.geometry.dofmap
    igi = np.array(mesh.geometry.input_global_indices, dtype=np.int64)
    global_geom_input = igi[x_dofmap]
    global_geom_owner = adios4dolfinx.utils.index_owner(mesh.comm, global_geom_input.reshape(-1), num_nodes_global)
    for i in range(shape[1]):
        arr_i = adios4dolfinx.comm_helpers.send_dofs_and_recv_values(global_geom_input.reshape(-1), global_geom_owner, mesh.comm, local_input_data[:,i], local_input_range[0])
        dof_pos = x_dofmap.reshape(-1)*shape[1]+i
        uh.x.array[dof_pos] = arr_i
    infile.close()
    return uh
infile = Path("idealized_LV_ref_0.h5")


for data in ["f0", "n0", "phi", "psi", "s0"]:
    uh = read_mesh_data(infile, mesh, data)
    with dolfinx.io.XDMFFile(mesh.comm, f"{data}.xdmf", "w") as xdmf:
        xdmf.write_mesh(mesh)
        xdmf.write_function(uh)


2 Likes