Write VTK HDF mesh from DOLFINx

script.py

# Author: Jørgen S. Dokken
# SPDX-License-Identifier:    MIT
from pathlib import Path


from mpi4py import MPI


import h5py


import dolfinx
import numpy as np


# Install h5py with mpi support
# HDF5_MPI="ON" HDF5_DIR="/usr/local" python3 -m pip install --break-system-packages --no-binary=h5py h5py -U --no-build-isolation
comm = MPI.COMM_WORLD
domain = dolfinx.mesh.create_unit_square(
   comm, 4, 4)
gdim = domain.geometry.dim
filename = Path("test.vtkhdf")


inf = h5py.File(filename, "w", driver="mpio", comm=comm)
hdf = inf.create_group(np.string_("VTKHDF"))


g_imap = domain.geometry.index_map()
num_nodes_local = g_imap.size_local
geom_range = g_imap.local_range
geom = domain.geometry.x[:num_nodes_local].copy()


dt = h5py.string_dtype(encoding="ascii")
hdf.attrs["Version"] = [2, 2]
hdf.attrs["Type"] = np.string_("UnstructuredGrid")


# Put local geometry
global_shape = (g_imap.size_global, 3)
geom_set = hdf.create_dataset(np.string_("Points"), global_shape, dtype=geom.dtype)
geom_set[geom_range[0]:geom_range[1], :] = geom






# Compute global topology
g_imap = g_imap
g_dmap = domain.geometry.dofmap
num_cells_local = domain.topology.index_map(domain.topology.dim).size_local
num_cells_global = domain.topology.index_map(domain.topology.dim).size_global
cell_range = domain.topology.index_map(domain.topology.dim).local_range
cmap = domain.geometry.cmap
geom_layout = cmap.create_dof_layout()
num_dofs_per_cell = geom_layout.num_entity_closure_dofs(domain.topology.dim)
dofs_out = np.zeros((num_cells_local, num_dofs_per_cell), dtype=np.int64)
assert g_dmap.shape[1] == num_dofs_per_cell
dofs_out[:, :] = np.asarray(
g_imap.local_to_global(g_dmap[:num_cells_local, :].reshape(-1))
).reshape(dofs_out.shape)
cell_type = domain.topology.cell_type
map_vtk = np.argsort(dolfinx.cpp.io.perm_vtk(cell_type, num_dofs_per_cell))
dofs_out = dofs_out[:, map_vtk]


# Put global topology
geom_set = hdf.create_dataset("Connectivity", (num_cells_global*num_dofs_per_cell,), dtype=np.int64)
geom_set[cell_range[0]*num_dofs_per_cell:num_dofs_per_cell*cell_range[1]] = dofs_out.reshape(-1)


# Put cell type
# Types
type_set = hdf.create_dataset("Types", (num_cells_global,), dtype=np.uint8)
cts = np.full(num_cells_local, dolfinx.cpp.io.get_vtk_cell_type(cell_type, domain.topology.dim))
type_set[cell_range[0]:cell_range[1]] = cts


# Geom dofmap offset
con_part = hdf.create_dataset("NumberOfConnectivityIds", (1,), dtype=np.int64)
if domain.comm.rank == 0:
   con_part[domain.comm.rank] = num_cells_global*num_dofs_per_cell


# Num cells
num_cells = hdf.create_dataset("NumberOfCells", (1,), dtype=np.int64)
if domain.comm.rank == 0:
   num_cells[domain.comm.rank] = num_cells_global


# num points
num_points = hdf.create_dataset("NumberOfPoints", (1,), dtype=np.int64)
if domain.comm.rank == 0:
   num_points[domain.comm.rank] = g_imap.size_global


# Offsets
offsets = hdf.create_dataset("Offsets", shape=(num_cells_global+1,), dtype=np.int64)
offsets[cell_range[0] + 1: cell_range[1] + 1] = np.arange(1, num_cells_local+1)*dofs_out.shape[1] + cell_range[0]*dofs_out.shape[1]
inf.close()