From 0313f06e33ddca43d196e4366534c8ce6384e125 Mon Sep 17 00:00:00 2001 From: mattfil Date: Wed, 11 Nov 2026 10:24:06 +0100 Subject: [PATCH] feat(vtkhdf compression): implemented compression, bmarks and docs. --- benchmarks/report_vtkhdf_compression.py | 137 ++++++++++++ benchmarks/vtkhdf_compression.py | 143 ++++++++++++ docs/backends/vtkhdf.rst | 87 ++++++++ src/io4dolfinx/backends/vtkhdf/backend.py | 141 +++++++++++- src/io4dolfinx/checkpointing.py | 8 +- tests/test_vtkhdf_compression.py | 251 ++++++++++++++++++++++ 6 files changed, 757 insertions(+), 10 deletions(-) create mode 100644 benchmarks/report_vtkhdf_compression.py create mode 100644 benchmarks/vtkhdf_compression.py create mode 100644 tests/test_vtkhdf_compression.py diff --git a/benchmarks/report_vtkhdf_compression.py b/benchmarks/report_vtkhdf_compression.py new file mode 100644 index 0000000..7a3ca38 --- /dev/null +++ b/benchmarks/report_vtkhdf_compression.py @@ -0,0 +1,137 @@ +"""Summarize rank subdirectories produced by vtkhdf_compression.py. + +Usage: python benchmarks/report_vtkhdf_compression.py output/vtkhdf-compression +""" + +import argparse +import csv +import json +import os +import statistics +from pathlib import Path + +os.environ.setdefault("MPLCONFIGDIR", "/tmp/io4dolfinx-matplotlib") + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +MODES = ["uncompressed", "gzip", "gzip_shuffle"] +LABELS = ["No compression", "Gzip level 4", "Gzip level 4 + shuffle"] +COLORS = ["#697586", "#3178b8", "#c45525"] + + +def summarize(root: Path): + reports = [json.loads(path.read_text()) for path in sorted(root.glob("ranks-*/results.json"))] + if not reports: + raise ValueError(f"No rank reports in {root}") + raw = [row for report in reports for row in report["results"]] + sizes = sorted({row["n"] for row in raw}) + ranks = sorted({row["ranks"] for row in raw}) + rows = [] + for n in sizes: + for rank in ranks: + for mode in MODES: + samples = [ + r for r in raw if r["n"] == n and r["ranks"] == rank and r["mode"] == mode + ] + if not samples: + raise ValueError(f"Missing n={n}, ranks={rank}, mode={mode}") + row = dict(n=n, cells=2 * n * n, ranks=rank, mode=mode, samples=len(samples)) + for metric in ( + "bytes", + "write_seconds", + "read_seconds", + "sync_seconds", + "mesh_write_seconds", + "point_write_seconds", + "cell_write_seconds", + "mesh_read_seconds", + "field_read_seconds", + ): + values = sorted(r[metric] for r in samples) + row[metric] = statistics.median(values) + if metric in ("write_seconds", "read_seconds"): + quartiles = statistics.quantiles(values, n=4, method="inclusive") + row[f"{metric}_p25"] = quartiles[0] + row[f"{metric}_p75"] = quartiles[2] + row[f"{metric}_min"] = min(values) + row[f"{metric}_max"] = max(values) + rows.append(row) + for row in rows: + baseline = next( + r + for r in rows + if r["n"] == row["n"] and r["ranks"] == row["ranks"] and r["mode"] == "uncompressed" + ) + row["size_ratio"] = row["bytes"] / baseline["bytes"] + row["write_time_ratio"] = row["write_seconds"] / baseline["write_seconds"] + row["read_time_ratio"] = row["read_seconds"] / baseline["read_seconds"] + with (root / "summary.csv").open("w", newline="") as stream: + writer = csv.DictWriter(stream, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + (root / "all_results.json").write_text(json.dumps(reports, indent=2) + "\n") + + plt.rcParams.update({"font.size": 10, "axes.spines.top": False, "axes.spines.right": False}) + fig, axes = plt.subplots(len(sizes), 3, figsize=(15, 3.5 * len(sizes)), squeeze=False) + metrics = [ + ("write_seconds", "Write time incl. fsync (ms)", 1000), + ("read_seconds", "Read time, warm cache (ms)", 1000), + ("bytes", "File size (MB, decimal)", 1e-6), + ] + for i, n in enumerate(sizes): + for j, (metric, label, scale) in enumerate(metrics): + ax = axes[i, j] + for mode, legend, color in zip(MODES, LABELS, COLORS): + selected = [ + next( + r for r in rows if r["n"] == n and r["ranks"] == rank and r["mode"] == mode + ) + for rank in ranks + ] + y = [r[metric] * scale for r in selected] + error = None + if metric != "bytes": + error = [ + [(r[metric] - r[f"{metric}_p25"]) * scale for r in selected], + [(r[f"{metric}_p75"] - r[metric]) * scale for r in selected], + ] + ax.errorbar(ranks, y, yerr=error, marker="o", capsize=3, color=color, label=legend) + ax.set_xscale("log", base=2) + ax.set_xticks(ranks, [str(rank) for rank in ranks]) + ax.set_xlabel("MPI ranks") + ax.set_ylabel(label) + ax.set_title(f"{2 * n * n:,} triangles ({n} × {n} mesh)") + ax.grid(True, alpha=0.2) + ax.set_ylim(bottom=0) + handles, labels = axes[0, 0].get_legend_handles_labels() + fig.legend(handles, labels, loc="upper center", ncol=3, bbox_to_anchor=(0.5, 0.965)) + fig.suptitle("VTKHDF: compression, shuffle, and MPI scaling", fontsize=17, y=0.995) + fig.text( + 0.5, + 0.005, + "Median of 5 measured runs; bars show the middle 50%. One warm-up per mode. " + "Smooth scalar/vector fields; automatic chunks; ext4 storage.", + ha="center", + fontsize=10, + ) + fig.tight_layout(rect=(0, 0.025, 1, 0.94)) + fig.savefig(root / "comparison.png", dpi=160) + fig.savefig(root / "comparison.svg") + plt.close(fig) + for row in rows: + print( + f"n={row['n']:3} ranks={row['ranks']} {row['mode']:14} " + f"{row['bytes'] / 1e6:7.3f} MB " + f"write={row['write_seconds'] * 1000:8.2f} ms " + f"read={row['read_seconds'] * 1000:8.2f} ms " + f"write/base={row['write_time_ratio']:.2f}" + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("root", type=Path) + summarize(parser.parse_args().root) diff --git a/benchmarks/vtkhdf_compression.py b/benchmarks/vtkhdf_compression.py new file mode 100644 index 0000000..964309c --- /dev/null +++ b/benchmarks/vtkhdf_compression.py @@ -0,0 +1,143 @@ +"""Compare uncompressed VTKHDF, gzip, and gzip+shuffle under MPI. + +Run from the repository root, for example: + mpiexec -n 2 python benchmarks/vtkhdf_compression.py --output /tmp/vtkhdf-mpi + +Uses smooth scalar/vector fields on triangular unit-square meshes. Timings include +open/close, mesh repartitioning, and field interpolation, and report the slowest +rank. Reads use warm filesystem caches; writes close the file without an explicit +fsync unless --sync is supplied. Results describe this workload and filesystem, +not general speed guarantees. +Each mode is warmed up before measurement; measured mode order rotates each round. +Output files are overwritten between repetitions to bound disk usage. +""" + +import argparse +import json +import os +import platform +from pathlib import Path + +from mpi4py import MPI + +import dolfinx +import h5py +import numpy as np +from dolfinx.fem import Function, functionspace +from dolfinx.mesh import create_unit_square + +import io4dolfinx + +MODES = { + "uncompressed": None, + "gzip": {"compression": "gzip", "compression_opts": 4, "shuffle": False}, + "gzip_shuffle": {"compression": "gzip", "compression_opts": 4, "shuffle": True}, +} + + +def measure(filename, mesh, point, cell, options, sync): + comm = mesh.comm + comm.barrier() + start = MPI.Wtime() + io4dolfinx.write_mesh(filename, mesh, backend="vtkhdf", backend_args=options) + mesh_written = MPI.Wtime() + io4dolfinx.write_point_data( + filename, point, time=0.0, mode=io4dolfinx.FileMode.append, backend_args=options + ) + point_written = MPI.Wtime() + io4dolfinx.write_cell_data( + filename, cell, time=0.0, mode=io4dolfinx.FileMode.append, backend_args=options + ) + cell_written = MPI.Wtime() + if sync: + comm.barrier() + if comm.rank == 0: + with filename.open("rb") as stream: + os.fsync(stream.fileno()) + comm.barrier() + write_finished = MPI.Wtime() + write_times = { + "write_seconds": write_finished - start, + "mesh_write_seconds": mesh_written - start, + "point_write_seconds": point_written - mesh_written, + "cell_write_seconds": cell_written - point_written, + "sync_seconds": write_finished - cell_written, + } + # Reduce outside the timed region; retain the critical-path time across ranks. + timings = {key: comm.allreduce(value, op=MPI.MAX) for key, value in write_times.items()} + comm.barrier() + start = MPI.Wtime() + restored = io4dolfinx.read_mesh(filename, comm, backend="vtkhdf") + mesh_read = MPI.Wtime() + io4dolfinx.read_point_data(filename, point.name, restored, 0.0, backend="vtkhdf") + io4dolfinx.read_cell_data(filename, cell.name, restored, 0.0, backend="vtkhdf") + fields_read = MPI.Wtime() + read_times = { + "read_seconds": fields_read - start, + "mesh_read_seconds": mesh_read - start, + "field_read_seconds": fields_read - mesh_read, + } + timings.update({key: comm.allreduce(value, op=MPI.MAX) for key, value in read_times.items()}) + return timings + + +def benchmark(output: Path, sizes: list[int], repeats: int, warmups: int, sync: bool): + comm = MPI.COMM_WORLD + if comm.rank == 0: + output.mkdir(parents=True, exist_ok=True) + comm.barrier() + results = [] + for n in sizes: + mesh = create_unit_square(comm, n, n) + point = Function(functionspace(mesh, ("Lagrange", 1, (3,))), name="velocity") + point.interpolate(lambda x: np.vstack((np.sin(x[0]), np.cos(x[1]), x[0] * x[1]))) + cell = Function(functionspace(mesh, ("DG", 0)), name="pressure") + cell.interpolate(lambda x: x[0] + 2 * x[1]) + for repeat in range(-warmups, repeats): + modes = list(MODES) + shift = max(repeat, 0) % len(modes) + modes = modes[shift:] + modes[:shift] + for mode in modes: + filename = output / f"n{n}-{mode}.vtkhdf" + timings = measure(filename, mesh, point, cell, MODES[mode], sync) + if repeat >= 0 and comm.rank == 0: + row = dict( + n=n, + cells=2 * n * n, + ranks=comm.size, + mode=mode, + repeat=repeat, + bytes=filename.stat().st_size, + **timings, + ) + results.append(row) + print(json.dumps(row), flush=True) + if comm.rank == 0: + # Save incrementally so completed sizes survive an interrupted run. + report = dict( + hostname=platform.node(), + platform=platform.platform(), + h5py=h5py.__version__, + hdf5=h5py.version.hdf5_version, + dolfinx=dolfinx.__version__, + mpi=MPI.Get_library_version().strip(), + repeats=repeats, + warmups=warmups, + fsync=sync, + modes=MODES, + results=results, + ) + (output / "results.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--sizes", type=int, nargs="+", default=[16, 256, 768]) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--warmups", type=int, default=1) + parser.add_argument("--sync", action="store_true", help="Include fsync after all writes") + args = parser.parse_args() + if any(n <= 0 for n in args.sizes) or args.repeats <= 0 or args.warmups < 0: + parser.error("sizes and repeats must be positive; warmups must be nonnegative") + benchmark(args.output, args.sizes, args.repeats, args.warmups, args.sync) diff --git a/docs/backends/vtkhdf.rst b/docs/backends/vtkhdf.rst index be309b6..ab78154 100644 --- a/docs/backends/vtkhdf.rst +++ b/docs/backends/vtkhdf.rst @@ -6,6 +6,93 @@ VTKHDF-backend Relies on MPI compatible :code:`h5py`, similar to the :ref:`H5PY backend`. See that backend for detailed installation instructions. +Lossless compression +~~~~~~~~~~~~~~~~~~~~ + +Compression is opt-in through ``backend_args``. The public writer signatures are +unchanged. Reuse the same dictionary on each call to compress both the mesh and +its point or cell data: + +.. code-block:: python + + from mpi4py import MPI + from dolfinx.fem import Function, functionspace + from dolfinx.mesh import create_unit_square + import io4dolfinx + + mesh = create_unit_square(MPI.COMM_WORLD, 16, 16) + u = Function(functionspace(mesh, ("Lagrange", 1)), name="u") + u.interpolate(lambda x: x[0] + x[1]) + q = Function(functionspace(mesh, ("DG", 0)), name="q") + q.interpolate(lambda x: x[0] - x[1]) + + storage_args = { + "compression": "gzip", + "compression_opts": 4, + "shuffle": True, + } + io4dolfinx.write_mesh( + "results.vtkhdf", mesh, backend="vtkhdf", backend_args=storage_args + ) + io4dolfinx.write_point_data( + "results.vtkhdf", u, time=0.0, mode=io4dolfinx.FileMode.append, + backend="vtkhdf", backend_args=storage_args, + ) + io4dolfinx.write_cell_data( + "results.vtkhdf", q, time=0.0, mode=io4dolfinx.FileMode.append, + backend="vtkhdf", backend_args=storage_args, + ) + +Supported options are ``compression`` (``"gzip"`` or ``None``), +``compression_opts`` (an integer from 0 to 9, requiring ``compression="gzip"``; +the gzip default is 4), and ``shuffle`` (a boolean, default ``False``). +Shuffle is lossless and can improve gzip compression. Other compression filters, +lossy scale-offset, checksums, and explicit chunk shapes are not supported by this +interface. Omitting these options creates uncompressed datasets. + +Coordinates, connectivity, topology offsets, cell types, and point/cell values +are compressed independently. Small bookkeeping arrays, including timestep +metadata, remain uncompressed. Each dataset uses h5py automatic chunking; +``chunks=True`` is accepted, but a fixed chunk shape is not. A chunk contains +elements of one dataset, not a combination of mesh and field arrays. See the +`h5py chunking guidance `_. + +Settings apply when a dataset is first created. On append, omitted options retain +that dataset's existing settings; explicitly conflicting settings raise +``ValueError`` before datasets are resized. Newly created fields use the options +of their own write call and do not inherit settings from ``write_mesh``. +Changing an existing dataset's filters requires rewriting it to a new file. + +MPI writes to filtered datasets are collective automatically. Every rank in the +communicator must call the writer, including ranks with no local data, and must +pass identical filter options. Parallel compression requires MPI-enabled h5py +and HDF5 with parallel filtered-write support (introduced in HDF5 1.10.2). +Invalid options and unavailable filters raise errors rather than silently +disabling compression. Readers do not need compression arguments. + +Compression reduces storage for compressible data, but CPU work, chunk sharing, +and MPI coordination can increase write time, particularly for small datasets. +There is no automatic size threshold or speed guarantee. Benchmark representative +data and rank counts; the repository's ``benchmarks/vtkhdf_compression.py`` reports +file sizes and mesh/field write and read times for uncompressed output, gzip, +and gzip with shuffle. For example, from the repository root: + +.. code-block:: bash + + OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 mpiexec -n 4 \ + python benchmarks/vtkhdf_compression.py \ + --output output/compression/ranks-4 --sizes 16 256 768 --repeats 5 --sync + +Each mode is warmed up before measurement and the measured order rotates between +repetitions. ``--sync`` includes a final filesystem ``fsync`` in write time; +reads use warm caches. The workload contains smooth scalar/vector fields, so +compression ratios do not predict those of noisy data. To combine runs in +``ranks-*`` directories into CSV and charts, run +``python benchmarks/report_vtkhdf_compression.py output/compression`` in an +environment with matplotlib. Compression enables collective writes, so timing +differences include the change of I/O mode as well as the filter itself. See the +`HDF5 parallel compression guidance `_. + .. automodule:: io4dolfinx.backends.vtkhdf.backend :members: :exclude-members: read_attributes, read_timestamps, write_attributes, write_mesh, write_meshtags, read_meshtags_data, read_dofmap, read_dofs, read_cell_perms, write_function, read_legacy_mesh, snapshot_checkpoint, read_hdf5_array diff --git a/src/io4dolfinx/backends/vtkhdf/backend.py b/src/io4dolfinx/backends/vtkhdf/backend.py index ddb464e..03448f4 100644 --- a/src/io4dolfinx/backends/vtkhdf/backend.py +++ b/src/io4dolfinx/backends/vtkhdf/backend.py @@ -2,6 +2,7 @@ Module that can read the VTKHDF format using h5py. """ +from contextlib import closing from pathlib import Path from typing import Any @@ -485,6 +486,100 @@ def read_function_names( return list(function_names) +def _validate_storage_options(comm: MPI.Comm, arguments: dict[str, Any]) -> dict[str, Any]: + """Extract explicit filter settings and check local HDF5 capabilities.""" + options: dict[str, Any] = { + key: arguments[key] + for key in ("compression", "compression_opts", "shuffle") + if key in arguments + } + compression = options.get("compression") + if compression is not None and compression != "gzip": + raise ValueError("VTKHDF compression must be 'gzip' or None") + if "compression_opts" in options: + level = options["compression_opts"] + if compression != "gzip" or type(level) is not int or not 0 <= level <= 9: + raise ValueError("compression_opts requires gzip and an integer level from 0 to 9") + if "shuffle" in options and type(options["shuffle"]) is not bool: + raise ValueError("shuffle must be a boolean") + if "chunks" in arguments and arguments["chunks"] is not True: + raise ValueError("VTKHDF uses automatic chunking; only chunks=True is supported") + if any(key in arguments for key in ("scaleoffset", "fletcher32")): + raise ValueError("VTKHDF supports only gzip compression and optional shuffle") + if compression is not None or options.get("shuffle", False): + parallel_hdf5 = h5py.get_config().mpi + if (comm.size > 1 and not parallel_hdf5) or ( + parallel_hdf5 + and (h5py.version.hdf5_version_tuple < (1, 10, 2) or MPI.Get_version() < (3, 0)) + ): + raise ValueError("Parallel VTKHDF compression requires MPI-enabled HDF5 >= 1.10.2") + filters = [h5py.h5z.FILTER_DEFLATE] if compression else [] + if options.get("shuffle", False): + filters.append(h5py.h5z.FILTER_SHUFFLE) + for filter_id in filters: + if not h5py.h5z.filter_avail(filter_id) or not ( + h5py.h5z.get_filter_info(filter_id) & h5py.h5z.FILTER_CONFIG_ENCODE_ENABLED + ): + raise ValueError(f"HDF5 filter {filter_id} is unavailable for writing") + return options + + +def _storage_options(comm: MPI.Comm, arguments: dict[str, Any]) -> dict[str, Any]: + """Agree on filter settings and validation errors before opening a file.""" + options: dict[str, Any] = {} + error = None + try: + options = _validate_storage_options(comm, arguments) + except (TypeError, ValueError) as exc: + error = str(exc) + + # One exchange handles both errors and settings; a single rank needs no exchange. + settings = comm.allgather((error, options)) if comm.size > 1 else [(error, options)] + for rank_error, _ in settings: + if rank_error is not None: + raise ValueError(rank_error) + if any(rank_options != options for _, rank_options in settings): + raise ValueError("VTKHDF compression options must be identical on every MPI rank") + return options + + +def _check_storage_options(dataset: h5py.Dataset, options: dict[str, Any]) -> None: + """Reject explicit changes to the immutable filters of an existing dataset.""" + for key, requested in options.items(): + actual = getattr(dataset, key) + if actual != requested: + raise ValueError( + f"Dataset {dataset.name} already has {key}={actual!r}; " + f"cannot append with {key}={requested!r}" + ) + + +def _write_dataset( + dataset: h5py.Dataset, selection: slice | tuple[slice, ...], values: np.ndarray +) -> None: + """Write a local slice, keeping empty ranks in filtered MPI writes.""" + if not (dataset.compression or dataset.shuffle) or dataset.file.driver != "mpio": + dataset[selection] = values + return + + if values.size: + with dataset.collective: + dataset[selection] = values + return + + # h5py's high-level assignment skips empty selections entirely. HDF5 still + # requires an H5Dwrite from these ranks for collective I/O. + with ( + closing(dataset.id.get_space()) as file_space, + closing(h5py.h5s.create_simple(values.shape)) as memory_space, + closing(h5py.h5p.create(h5py.h5p.DATASET_XFER)) as transfer, + ): + file_space.select_none() + memory_space.select_none() + transfer.set_dxpl_mpio(h5py.h5fd.MPIO_COLLECTIVE) + dataset.id.write(memory_space, file_space, values, dxpl=transfer) + + def _create_dataset( root: h5py.File | h5py.Group, name: str, @@ -494,12 +589,18 @@ def _create_dataset( maxshape: tuple[int | None, ...], mode: str, resize: bool = True, + storage_options: dict[str, Any] | None = None, ) -> h5py.Dataset: if name not in root.keys(): mode = "w" if mode == "w": dataset = root.create_dataset( - name, shape=shape, dtype=dtype, chunks=chunks, maxshape=maxshape + name, + shape=shape, + dtype=dtype, + chunks=chunks, + maxshape=maxshape, + **(storage_options or {}), ) elif mode == "a": dataset = root[name] @@ -559,8 +660,15 @@ def write_mesh( h5_mode = convert_file_mode(mode) backend_args = get_default_backend_args(backend_args) name = backend_args["name"] + storage_options = _storage_options(comm, backend_args) with h5pyfile(filename, h5_mode, comm=comm) as h5file: + # Check all existing bulk arrays before changing attributes or resizing anything. + if storage_options and h5_mode == "a" and f"/VTKHDF/{name}" in h5file: + block = h5file[f"/VTKHDF/{name}"] + for key in ("Points", "Offsets", "Connectivity", "Types"): + if key in block: + _check_storage_options(block[key], storage_options) hdf = _create_group(h5file, "/VTKHDF", h5_mode) hdf.attrs.create("Type", "MultiBlockDataSet") hdf.attrs["Version"] = _vtk_hdf_version @@ -605,11 +713,14 @@ def write_mesh( chunks=True, maxshape=(None, 3), mode=h5_mode, + storage_options=storage_options, ) insert_slice = _compute_append_slice( points, mesh.num_nodes_global, mesh.local_geometry_pos, h5_mode ) - points[insert_slice, : mesh.local_geometry.shape[1]] = mesh.local_geometry + _write_dataset( + points, (insert_slice, slice(0, mesh.local_geometry.shape[1])), mesh.local_geometry + ) # NOTE: VTKHDF currently does not support reading time dependent topology in # Paraview: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 @@ -647,16 +758,21 @@ def write_mesh( mode=h5_mode, maxshape=(None,), resize=False, # Resize should really be True + storage_options=storage_options, ) - offset_data = np.arange(0, mesh.local_topology.size + 1, mesh.local_topology.shape[1]) + # Each offset has one writer. Only the last nonempty rank writes the final offset. + first, last = mesh.local_topology_pos + include_end = last == mesh.num_cells_global and last > first + offset_data = np.arange(last - first + int(include_end), dtype=np.int64) + offset_data *= num_dofs_per_cell offset_data += num_dofs_per_cell * mesh.local_topology_pos[0] insert_slice = _compute_append_slice( offsets, mesh.num_cells_global + 1, - (mesh.local_topology_pos[0], mesh.local_topology_pos[1] + 1), + (first, last + int(include_end)), mode=h5_mode, ) - offsets[insert_slice] = offset_data + _write_dataset(offsets, insert_slice, offset_data) del offset_data # Permute and store topology data @@ -672,6 +788,7 @@ def write_mesh( maxshape=(None,), mode=h5_mode, resize=False, # Resize should really be True, see issue below + storage_options=storage_options, ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 insert_slice = _compute_append_slice( topology, @@ -679,13 +796,14 @@ def write_mesh( np.array(mesh.local_topology_pos) * num_dofs_per_cell, mode=h5_mode, ) - topology[insert_slice] = topology_data + _write_dataset(topology, insert_slice, topology_data) del topology_data # Store celltypes cell_types = np.full( mesh.local_topology.shape[0], fenicsx_compat.vtk_cell_type(dx_ct, num_dofs_per_cell), + dtype=np.uint8, ) types = _create_dataset( mesh_group, @@ -696,11 +814,12 @@ def write_mesh( chunks=True, mode=h5_mode, resize=False, # Resize should really be True, see issue below + storage_options=storage_options, ) # VTKHDFReader issue: https://gitlab.kitware.com/vtk/vtk/-/issues/19257 insert_slice = _compute_append_slice( types, mesh.num_cells_global, mesh.local_topology_pos, h5_mode ) - types[insert_slice] = cell_types + _write_dataset(types, insert_slice, cell_types) del cell_types steps = _create_group(mesh_group, "Steps", mode=h5_mode) @@ -1202,6 +1321,7 @@ def write_data( backend_args = get_default_backend_args(backend_args) mesh_name = backend_args["name"] extension = array_data.type + storage_options = _storage_options(comm, backend_args) with h5pyfile(filename, h5_mode, comm=comm) as h5file: hdf = h5file["VTKHDF"] file_type = _decode_bytes_if_needed(hdf.attrs["Type"]) @@ -1212,6 +1332,10 @@ def write_data( # Find mesh block to add data to block = _get_vtk_group(h5file, mesh_name) + data_path = f"{extension}Data/{array_data.name}" + if storage_options and data_path in block: + _check_storage_options(block[data_path], storage_options) + data_group = _create_group(block, f"{extension}Data", mode=h5_mode) dataset = _create_dataset( data_group, @@ -1221,6 +1345,7 @@ def write_data( chunks=True, maxshape=(None, array_data.values.shape[1]), mode=h5_mode, + storage_options=storage_options, ) insert_slice = _compute_append_slice( dataset, @@ -1233,7 +1358,7 @@ def write_data( ), mode=h5_mode, ) - dataset[insert_slice] = array_data.values + _write_dataset(dataset, insert_slice, array_data.values) steps = _create_group(block, "Steps", mode=h5_mode) pdo = _create_group(steps, f"{extension}DataOffsets", mode=h5_mode) diff --git a/src/io4dolfinx/checkpointing.py b/src/io4dolfinx/checkpointing.py index 1822873..246b4f8 100644 --- a/src/io4dolfinx/checkpointing.py +++ b/src/io4dolfinx/checkpointing.py @@ -577,6 +577,8 @@ def write_mesh( mesh: The mesh to write to file store_partition_info: Store mesh partitioning (including ghosting) to file + backend_args: Backend-specific options. For VTKHDF compression, + see :ref:`vtkhdfbackend`. mesh_name: Name of the mesh in the file to act on. Defaults to the mesh written without an explicit name, so existing checkpoints and call sites are unaffected. @@ -733,7 +735,8 @@ def write_point_data( u: The function to store time: Time stamp mode: Append or write - backend_args: The backend arguments + backend_args: Backend-specific options. For VTKHDF compression, + see :ref:`vtkhdfbackend`. backend: Which backend to use. mesh_name: Name of the mesh in the file to write the data against. Defaults to the mesh written without an explicit name. @@ -777,7 +780,8 @@ def write_cell_data( point_data: Data to write to file time: Time stamp mode: Append or write - backend_args: The backend arguments + backend_args: Backend-specific options. For VTKHDF compression, + see :ref:`vtkhdfbackend`. mesh_name: Name of the mesh in the file to write the data against. Defaults to the mesh written without an explicit name. """ diff --git a/tests/test_vtkhdf_compression.py b/tests/test_vtkhdf_compression.py new file mode 100644 index 0000000..9a5ed05 --- /dev/null +++ b/tests/test_vtkhdf_compression.py @@ -0,0 +1,251 @@ +"""Lossless VTKHDF filters, append policies, and collective empty selections.""" + +from dataclasses import replace + +from mpi4py import MPI + +import h5py +import numpy as np +import pytest +from dolfinx.fem import Function, functionspace +from dolfinx.mesh import create_unit_square +from vtkmodules.vtkIOHDF import vtkHDFReader + +import io4dolfinx +from io4dolfinx.backends.vtkhdf import backend +from io4dolfinx.structures import ArrayData +from io4dolfinx.writers import prepare_meshdata_for_storage + + +@pytest.mark.parametrize("compression", [None, "gzip"]) +@pytest.mark.parametrize("shuffle", [False, True]) +def test_compressed_roundtrip(shared_tmp_path, compression, shuffle): + comm = MPI.COMM_WORLD + filename = shared_tmp_path / "roundtrip.vtkhdf" + mesh = create_unit_square(comm, 5, 4) + options = {"compression": compression, "shuffle": shuffle} + if compression: + options["compression_opts"] = 4 + io4dolfinx.write_mesh(filename, mesh, backend="vtkhdf", backend_args=options) + for time in (0.0, 1.0): + # Omitting options on append must retain the existing mesh filters. + if time: + io4dolfinx.write_mesh( + filename, mesh, time=time, mode=io4dolfinx.FileMode.append, backend="vtkhdf" + ) + for association in ("point", "cell"): + for components in (1, 3): + family, degree = ("Lagrange", 1) if association == "point" else ("DG", 0) + shape = () if components == 1 else (3,) + V = functionspace(mesh, (family, degree, shape)) + u = Function(V, name=f"{association}{components}") + u.interpolate( + (lambda x: x[0] + time) + if components == 1 + else (lambda x: np.vstack((x[0] + time, x[1], 2 * x[0]))) + ) + writer = getattr(io4dolfinx, f"write_{association}_data") + writer( + filename, + u, + time=time, + mode=io4dolfinx.FileMode.append, + backend_args=options if time == 0.0 else None, + ) + reader = getattr(io4dolfinx, f"read_{association}_data") + restored = io4dolfinx.read_mesh(filename, comm, time=time, backend="vtkhdf") + result = reader(filename, mesh=restored, name=u.name, time=time, backend="vtkhdf") + expected = Function(result.function_space) + expected.interpolate( + (lambda x: x[0] + time) + if components == 1 + else (lambda x: np.vstack((x[0] + time, x[1], 2 * x[0]))) + ) + np.testing.assert_allclose(result.x.array, expected.x.array, atol=1e-14) + + # A newly created field does not inherit the mesh's filters. + plain = Function(functionspace(mesh, ("Lagrange", 1)), name="plain") + io4dolfinx.write_point_data( + filename, plain, time=1.0, mode=io4dolfinx.FileMode.append, backend_args=None + ) + with h5py.File(filename, "r") as hdf: + block = hdf["VTKHDF/mesh"] + for path in ( + "Points", + "Connectivity", + "Offsets", + "Types", + "PointData/point1", + "PointData/point3", + "CellData/cell1", + "CellData/cell3", + ): + dataset = block[path] + assert dataset.compression == compression + assert dataset.shuffle == shuffle + assert dataset.chunks is not None + if compression: + assert dataset.compression_opts == 4 + for path in ("NumberOfPoints", "NumberOfCells", "Steps/Values", "PointData/plain"): + assert block[path].compression is None + assert not block[path].shuffle + + for time in (0.0, 1.0): + restored = io4dolfinx.read_mesh(filename, comm, time=time, backend="vtkhdf") + assert restored.topology.index_map(2).size_global == mesh.topology.index_map(2).size_global + restored_points = np.vstack( + comm.allgather(restored.geometry.x[: restored.geometry.index_map().size_local]) + ) + original_points = np.vstack( + comm.allgather(mesh.geometry.x[: mesh.geometry.index_map().size_local]) + ) + np.testing.assert_allclose( + np.sort(restored_points, axis=0), np.sort(original_points, axis=0) + ) + vtk_reader = vtkHDFReader() + vtk_reader.SetFileName(str(filename)) + vtk_reader.UpdateTimeStep(time) + output = vtk_reader.GetOutputDataObject(0) + while output.IsA("vtkMultiBlockDataSet"): + output = output.GetBlock(0) + assert output.GetNumberOfPoints() == mesh.geometry.index_map().size_global + assert output.GetNumberOfCells() == mesh.topology.index_map(2).size_global + for association in ("point", "cell"): + attributes = output.GetPointData() if association == "point" else output.GetCellData() + for components in (1, 3): + values = attributes.GetArray(f"{association}{components}") + assert values is not None + assert values.GetNumberOfComponents() == components + # Scalar values are x + time at points, or at cell centroids. + for i in range(values.GetNumberOfTuples()): + if association == "point": + x = output.GetPoint(i)[0] + else: + cell = output.GetCell(i) + x = np.mean( + [ + output.GetPoint(cell.GetPointId(j))[0] + for j in range(cell.GetNumberOfPoints()) + ] + ) + assert np.isclose(values.GetComponent(i, 0), x + time) + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("association", ["Point", "Cell"]) +def test_compressed_empty_ranks(shared_tmp_path, dtype, association): + """Empty ranks participate, and filters preserve special floating-point values.""" + comm = MPI.COMM_WORLD + filename = shared_tmp_path / "empty_ranks.vtkhdf" + mesh = create_unit_square(MPI.COMM_SELF, 2, 2, dtype=dtype) + data = prepare_meshdata_for_storage(mesh, False) + if comm.rank: + data = replace( + data, + local_geometry=data.local_geometry[:0], + local_topology=data.local_topology[:0], + local_geometry_pos=(data.num_nodes_global, data.num_nodes_global), + local_topology_pos=(data.num_cells_global, data.num_cells_global), + ) + options = {"compression": "gzip", "shuffle": True} + backend.write_mesh(filename, comm, data, options, io4dolfinx.FileMode.write, 0.0) + count = data.num_nodes_global if association == "Point" else data.num_cells_global + local_range = data.local_geometry_pos if association == "Point" else data.local_topology_pos + reference = np.resize(np.array([-0.0, np.nan, np.inf, -np.inf, 1.2345], dtype=dtype), count) + values = reference[slice(*local_range)].reshape(-1, 1).copy() + field = ArrayData("special", values, (count, 1), local_range, association) + backend.write_data(filename, field, comm, 0.0, io4dolfinx.FileMode.append, options) + with h5py.File(filename, "r") as hdf: + block = hdf["VTKHDF/mesh"] + np.testing.assert_array_equal(block["Offsets"][:], np.arange(0, 25, 3)) + stored = block[f"{association}Data/special"][:, 0] + # Compare bytes as well: equality alone does not distinguish +0 from -0. + assert stored.tobytes() == reference.tobytes() + + +@pytest.mark.parametrize( + "options", + [ + {"compression": "lzf"}, + {"compression": "gzip", "compression_opts": 10}, + {"compression": "gzip", "compression_opts": True}, + {"compression_opts": 4}, + {"shuffle": "yes"}, + {"chunks": (25000,)}, + {"scaleoffset": 2}, + ], +) +def test_invalid_options_do_not_create_file(shared_tmp_path, options): + comm = MPI.COMM_WORLD + mesh = create_unit_square(comm, 2, 2) + filename = shared_tmp_path / "invalid.vtkhdf" + with pytest.raises(ValueError): + io4dolfinx.write_mesh(filename, mesh, backend="vtkhdf", backend_args=options) + assert not filename.exists() + + +@pytest.mark.parametrize("target", ["mesh", "point", "cell"]) +@pytest.mark.parametrize( + "conflict", + [ + {"compression": None}, + {"compression": "gzip", "compression_opts": 1}, + {"shuffle": False}, + ], +) +def test_conflicting_append_preserves_file(shared_tmp_path, target, conflict): + comm = MPI.COMM_WORLD + filename = shared_tmp_path / "conflict.vtkhdf" + mesh = create_unit_square(comm, 2, 2) + options = {"compression": "gzip", "compression_opts": 4, "shuffle": True} + io4dolfinx.write_mesh(filename, mesh, backend="vtkhdf", backend_args=options) + u = Function(functionspace(mesh, ("Lagrange", 1)), name="u") + if target != "mesh": + getattr(io4dolfinx, f"write_{target}_data")( + filename, u, time=0.0, mode=io4dolfinx.FileMode.append, backend_args=options + ) + before = filename.read_bytes() + with pytest.raises(ValueError, match="cannot append"): + if target == "mesh": + io4dolfinx.write_mesh( + filename, + mesh, + time=1.0, + mode=io4dolfinx.FileMode.append, + backend="vtkhdf", + backend_args=conflict, + ) + else: + getattr(io4dolfinx, f"write_{target}_data")( + filename, u, time=1.0, mode=io4dolfinx.FileMode.append, backend_args=conflict + ) + assert filename.read_bytes() == before + + +def test_rank_local_filter_failure(shared_tmp_path, monkeypatch): + comm = MPI.COMM_WORLD + if comm.rank == 0: + monkeypatch.setattr(h5py.h5z, "filter_avail", lambda _: False) + mesh = create_unit_square(comm, 2, 2) + filename = shared_tmp_path / "unavailable.vtkhdf" + with pytest.raises(ValueError, match="unavailable"): + io4dolfinx.write_mesh( + filename, mesh, backend="vtkhdf", backend_args={"compression": "gzip"} + ) + assert not filename.exists() + + +def test_inconsistent_rank_options(shared_tmp_path): + comm = MPI.COMM_WORLD + if comm.size == 1: + pytest.skip("Requires multiple MPI ranks") + mesh = create_unit_square(comm, 2, 2) + filename = shared_tmp_path / "inconsistent.vtkhdf" + with pytest.raises(ValueError, match="identical"): + io4dolfinx.write_mesh( + filename, + mesh, + backend="vtkhdf", + backend_args={"compression": "gzip", "compression_opts": comm.rank % 10}, + ) + assert not filename.exists()