From 6cb56b9c92b36b0665977c22d35dc52fe569a224 Mon Sep 17 00:00:00 2001 From: Sammy Dabbas Date: Sun, 5 Jul 2026 01:41:57 -0400 Subject: [PATCH] Add Particle.write_vtkhdf to visualize lost particles in ParaView Adds a static method that reads one or more lost-particle restart files and writes a single VTKHDF PolyData file, one vertex per particle, with point data for energy, weight, direction, particle type, batch, generation, and id. Only h5py and numpy are required. See #3922. --- openmc/particle_restart.py | 92 +++++++++++++ tests/unit_tests/test_particle_restart.py | 156 ++++++++++++++++++++++ 2 files changed, 248 insertions(+) create mode 100644 tests/unit_tests/test_particle_restart.py diff --git a/openmc/particle_restart.py b/openmc/particle_restart.py index 44864107d9a..33ecc5db668 100644 --- a/openmc/particle_restart.py +++ b/openmc/particle_restart.py @@ -1,6 +1,11 @@ +from collections.abc import Iterable +import os + import h5py +import numpy as np import openmc.checkvalue as cv +from .checkvalue import PathLike from .particle_type import ParticleType _VERSION_PARTICLE_RESTART = 2 @@ -59,3 +64,90 @@ def __init__(self, filename): self.uvw = f['uvw'][()] self.weight = f['weight'][()] self.xyz = f['xyz'][()] + + @staticmethod + def write_vtkhdf( + particle_files: PathLike | Iterable[PathLike], + filename: PathLike = 'lost_particles.vtkhdf' + ): + """Write lost particle files to a single VTK HDF PolyData file. + + Each lost particle is written as a vertex with point data arrays for + energy, weight, direction, particle type, batch, generation, and + particle ID. The resulting file can be opened directly in ParaView to + visualize where particles are being lost. Only h5py and numpy are + required to write the file. + + .. versionadded:: 0.15.4 + + Parameters + ---------- + particle_files : path-like or iterable of path-like + Path(s) to lost particle restart files + filename : path-like + Name of the VTK HDF file to write + + Examples + -------- + >>> openmc.Particle.write_vtkhdf( + ... ['particle_10_542.h5', 'particle_12_178.h5'], + ... 'lost_particles.vtkhdf' + ... ) + + """ + if isinstance(particle_files, (str, os.PathLike)): + particle_files = [particle_files] + particles = [Particle(f) for f in particle_files] + if not particles: + raise ValueError('At least one particle file must be provided.') + + n = len(particles) + points = np.array([p.xyz for p in particles], dtype=np.float64) + + with h5py.File(filename, 'w') as f: + root = f.create_group('VTKHDF') + root.attrs['Version'] = (2, 1) + ascii_type = 'PolyData'.encode('ascii') + root.attrs.create( + 'Type', ascii_type, + dtype=h5py.string_dtype('ascii', len(ascii_type))) + + root.create_dataset('NumberOfPoints', data=(n,), dtype='i8') + root.create_dataset('Points', data=points) + + # Represent each lost particle as a vertex cell + vertices = root.create_group('Vertices') + vertices.create_dataset('NumberOfCells', data=(n,), dtype='i8') + vertices.create_dataset( + 'NumberOfConnectivityIds', data=(n,), dtype='i8') + vertices.create_dataset( + 'Connectivity', data=np.arange(n, dtype=np.int64)) + vertices.create_dataset( + 'Offsets', data=np.arange(n + 1, dtype=np.int64)) + + # The remaining topologies required by VTK HDF PolyData are empty + for group_name in ('Lines', 'Polygons', 'Strips'): + group = root.create_group(group_name) + group.create_dataset('NumberOfCells', data=(0,), dtype='i8') + group.create_dataset( + 'NumberOfConnectivityIds', data=(0,), dtype='i8') + group.create_dataset( + 'Connectivity', data=np.empty(0, dtype=np.int64)) + group.create_dataset( + 'Offsets', data=np.zeros(1, dtype=np.int64)) + + point_data = root.create_group('PointData') + point_data.create_dataset('energy', data=np.array( + [p.energy for p in particles], dtype=np.float64)) + point_data.create_dataset('weight', data=np.array( + [p.weight for p in particles], dtype=np.float64)) + point_data.create_dataset('direction', data=np.array( + [p.uvw for p in particles], dtype=np.float64)) + point_data.create_dataset('particle_type', data=np.array( + [int(p.type) for p in particles], dtype=np.int64)) + point_data.create_dataset('batch', data=np.array( + [p.current_batch for p in particles], dtype=np.int64)) + point_data.create_dataset('generation', data=np.array( + [p.current_generation for p in particles], dtype=np.int64)) + point_data.create_dataset('id', data=np.array( + [p.id for p in particles], dtype=np.int64)) diff --git a/tests/unit_tests/test_particle_restart.py b/tests/unit_tests/test_particle_restart.py new file mode 100644 index 00000000000..b5bd7b5820e --- /dev/null +++ b/tests/unit_tests/test_particle_restart.py @@ -0,0 +1,156 @@ +import h5py +import numpy as np +import pytest + +import openmc + + +def write_particle_restart_file(filename, *, current_batch=10, + current_generation=1, generations_per_batch=1, + n_particles=1000, run_mode='fixed source', + particle_id=542, particle_type=2112, + weight=1.0, energy=1.0e6, + xyz=(0., 0., 0.), uvw=(1., 0., 0.), + time=0.0): + """Create a file matching Particle::write_restart in src/particle.cpp.""" + with h5py.File(filename, 'w') as f: + f.attrs['filetype'] = np.bytes_('particle restart') + f.attrs['version'] = np.array([2, 1]) + f.create_dataset('current_batch', data=current_batch) + f.create_dataset('generations_per_batch', data=generations_per_batch) + f.create_dataset('current_generation', data=current_generation) + f.create_dataset('n_particles', data=n_particles) + f.create_dataset('run_mode', data=np.bytes_(run_mode)) + f.create_dataset('id', data=particle_id) + f.create_dataset('type', data=particle_type) + f.create_dataset('weight', data=weight) + f.create_dataset('energy', data=energy) + f.create_dataset('xyz', data=np.asarray(xyz)) + f.create_dataset('uvw', data=np.asarray(uvw)) + f.create_dataset('time', data=time) + + +@pytest.fixture +def particle_files(tmp_path): + params = [ + {'current_batch': 1, 'current_generation': 1, 'particle_id': 285880, + 'particle_type': 2112, 'weight': 1.0, 'energy': 2.0e6, + 'xyz': (1., 2., 3.), 'uvw': (1., 0., 0.)}, + {'current_batch': 1, 'current_generation': 2, 'particle_id': 300145, + 'particle_type': 22, 'weight': 0.5, 'energy': 5.0e5, + 'xyz': (-4., 5., -6.), 'uvw': (0., 1., 0.)}, + {'current_batch': 3, 'current_generation': 1, 'particle_id': 12, + 'particle_type': 2112, 'weight': 2.0, 'energy': 14.1e6, + 'xyz': (7., -8., 9.), 'uvw': (0., 0., -1.)}, + ] + files = [] + for i, kwargs in enumerate(params): + path = tmp_path / f'particle_{kwargs["current_batch"]}_{i}.h5' + write_particle_restart_file(path, **kwargs) + files.append(path) + return files, params + + +def test_read_particle_restart(particle_files): + files, params = particle_files + p = openmc.Particle(files[0]) + assert p.current_batch == params[0]['current_batch'] + assert p.current_generation == params[0]['current_generation'] + assert p.id == params[0]['particle_id'] + assert p.type == openmc.ParticleType.NEUTRON + assert p.weight == params[0]['weight'] + assert p.energy == params[0]['energy'] + assert p.run_mode == 'fixed source' + np.testing.assert_array_equal(p.xyz, params[0]['xyz']) + np.testing.assert_array_equal(p.uvw, params[0]['uvw']) + + +def test_write_vtkhdf(particle_files, tmp_path): + files, params = particle_files + out = tmp_path / 'lost_particles.vtkhdf' + openmc.Particle.write_vtkhdf(files, out) + + with h5py.File(out, 'r') as f: + root = f['VTKHDF'] + assert root.attrs['Type'] == b'PolyData' + assert root.attrs['Version'][0] >= 2 + + n = len(files) + assert root['NumberOfPoints'][()] == [n] + assert root['Points'].shape == (n, 3) + np.testing.assert_array_equal( + root['Points'][()], [p['xyz'] for p in params]) + + # Each point is a vertex cell + vertices = root['Vertices'] + assert vertices['NumberOfCells'][()] == [n] + assert vertices['NumberOfConnectivityIds'][()] == [n] + np.testing.assert_array_equal( + vertices['Connectivity'][()], np.arange(n)) + np.testing.assert_array_equal( + vertices['Offsets'][()], np.arange(n + 1)) + + # Other PolyData topologies present but empty + for name in ('Lines', 'Polygons', 'Strips'): + group = root[name] + assert group['NumberOfCells'][()] == [0] + assert group['NumberOfConnectivityIds'][()] == [0] + assert group['Connectivity'].shape == (0,) + np.testing.assert_array_equal(group['Offsets'][()], [0]) + + point_data = root['PointData'] + np.testing.assert_array_equal( + point_data['energy'][()], [p['energy'] for p in params]) + np.testing.assert_array_equal( + point_data['weight'][()], [p['weight'] for p in params]) + np.testing.assert_array_equal( + point_data['direction'][()], [p['uvw'] for p in params]) + np.testing.assert_array_equal( + point_data['particle_type'][()], + [p['particle_type'] for p in params]) + np.testing.assert_array_equal( + point_data['batch'][()], [p['current_batch'] for p in params]) + np.testing.assert_array_equal( + point_data['generation'][()], + [p['current_generation'] for p in params]) + np.testing.assert_array_equal( + point_data['id'][()], [p['particle_id'] for p in params]) + + +def test_write_vtkhdf_single_file(particle_files, tmp_path): + files, params = particle_files + out = tmp_path / 'single.vtkhdf' + openmc.Particle.write_vtkhdf(str(files[0]), out) + + with h5py.File(out, 'r') as f: + root = f['VTKHDF'] + assert root['NumberOfPoints'][()] == [1] + np.testing.assert_array_equal(root['Points'][()], [params[0]['xyz']]) + + +def test_write_vtkhdf_no_files(tmp_path): + with pytest.raises(ValueError): + openmc.Particle.write_vtkhdf([], tmp_path / 'empty.vtkhdf') + + +def test_write_vtkhdf_read_with_vtk(particle_files, tmp_path): + vtk = pytest.importorskip('vtk') + files, params = particle_files + out = tmp_path / 'lost_particles.vtkhdf' + openmc.Particle.write_vtkhdf(files, out) + + reader = vtk.vtkHDFReader() + reader.SetFileName(str(out)) + reader.Update() + data = reader.GetOutput() + + assert data.GetNumberOfPoints() == len(files) + assert data.GetNumberOfVerts() == len(files) + point_data = data.GetPointData() + names = {point_data.GetArrayName(i) + for i in range(point_data.GetNumberOfArrays())} + assert {'energy', 'weight', 'direction', 'particle_type', 'batch', + 'generation', 'id'} <= names + energy = point_data.GetArray('energy') + values = [energy.GetValue(i) for i in range(energy.GetNumberOfValues())] + assert values == [p['energy'] for p in params]