@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)