Reading node based data from "Legacy" XDMFFile

quick and "dirty" python version of reading data at geometry nodes. Should be properly implemented with ADIOS2.

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)

ref: https://fenicsproject.discourse.group/t/i-o-from-xdmf-hdf5-files-in-dolfin-x/3122/47

jorgensd / adios4dolfinx

Reading node based data from "Legacy" XDMFFile #16