From 365aa0b0eb187210666c1f946d4cbbd89b3f0289 Mon Sep 17 00:00:00 2001 From: Paul Romano Date: Tue, 17 May 2022 07:33:55 -0400 Subject: [PATCH] Write all tracks to a single file --- docs/source/pythonapi/base.rst | 1 + include/openmc/track_output.h | 17 +++++++++ openmc/trackfile.py | 65 ++++++++++++++++++++++++--------- scripts/openmc-track-to-vtk | 33 ++++++++--------- src/simulation.cpp | 10 ++++++ src/track_output.cpp | 66 +++++++++++++++++++--------------- 6 files changed, 130 insertions(+), 62 deletions(-) diff --git a/docs/source/pythonapi/base.rst b/docs/source/pythonapi/base.rst index fdd89b844..a54d2ff35 100644 --- a/docs/source/pythonapi/base.rst +++ b/docs/source/pythonapi/base.rst @@ -187,6 +187,7 @@ Post-processing openmc.Particle openmc.StatePoint openmc.Summary + openmc.Track openmc.TrackFile The following classes and functions are used for functional expansion reconstruction. diff --git a/include/openmc/track_output.h b/include/openmc/track_output.h index fc65a2eaa..c4641af2b 100644 --- a/include/openmc/track_output.h +++ b/include/openmc/track_output.h @@ -9,8 +9,25 @@ namespace openmc { // Non-member functions //============================================================================== +//! Open HDF5 track file for writing and create track datatype +void open_track_file(); + +//! Close HDF5 resources for track file +void close_track_file(); + +//! Create a new track state history for a primary/secondary particle +// +//! \param[in] p Current particle void add_particle_track(Particle& p); + +//! Store particle's current state +// +//! \param[in] p Current particle void write_particle_track(Particle& p); + +//! Write full particle state history to HDF5 track file +// +//! \param[in] p Current particle void finalize_particle_track(Particle& p); } // namespace openmc diff --git a/openmc/trackfile.py b/openmc/trackfile.py index ab2b81a54..4eaa72f4d 100644 --- a/openmc/trackfile.py +++ b/openmc/trackfile.py @@ -5,50 +5,57 @@ import h5py from .source import SourceParticle, ParticleType -TrackHistory = namedtuple('TrackHistory', ['particle', 'states']) +ParticleTrack = namedtuple('ParticleTrack', ['particle', 'states']) -class TrackFile: - """Particle track output +def _identifier(dset_name): + """Return (batch, gen, particle) tuple given dataset name""" + _, batch, gen, particle = dset_name.split('_') + return (int(batch), int(gen), int(particle)) + + +class Track: + """Tracks resulting from a single source particle Parameters ---------- - filepath : str or pathlib.Path - Path of file to load + dset : h5py.Dataset + Dataset to read track data from Attributes ---------- + identifier : tuple + Tuple of (batch, generation, particle number) + particles : list + List of tuples containing (particle type, array of track states) sources : list List of :class:`SourceParticle` representing each primary/secondary particle - tracks : list - List of tuples containing (particle type, array of track states) """ - def __init__(self, filepath): - # Read data from track file - with h5py.File(filepath, 'r') as fh: - offsets = fh.attrs['offsets'] - tracks = fh['tracks'][()] - particles = fh.attrs['particles'] + def __init__(self, dset): + tracks = dset[()] + offsets = dset.attrs['offsets'] + particles = dset.attrs['particles'] + self.identifier = _identifier(dset.name) # Construct list of track histories tracks_list = [] for particle, start, end in zip(particles, offsets[:-1], offsets[1:]): ptype = ParticleType(particle) - tracks_list.append(TrackHistory(ptype, tracks[start:end])) - self.tracks = tracks_list + tracks_list.append(ParticleTrack(ptype, tracks[start:end])) + self.particles = tracks_list def __repr__(self): - return f'' + return f'' def plot(self): """Produce a 3D plot of particle tracks""" import matplotlib.pyplot as plt fig = plt.figure() ax = plt.axes(projection='3d') - for _, states in self.tracks: + for _, states in self.particles: r = states['r'] ax.plot3D(r['x'], r['y'], r['z']) ax.set_xlabel('x [cm]') @@ -70,3 +77,27 @@ class TrackFile: ) ) return sources + + +class TrackFile(list): + """Collection of particle tracks + + Parameters + ---------- + filepath : str or pathlib.Path + Path of file to load + + Attributes + ---------- + tracks : list + List of :class:`Track` objects + + """ + + def __init__(self, filepath): + # Read data from track file + with h5py.File(filepath, 'r') as fh: + + for dset_name in sorted(fh, key=_identifier): + dset = fh[dset_name] + self.append(Track(dset)) diff --git a/scripts/openmc-track-to-vtk b/scripts/openmc-track-to-vtk index d1b2f828a..7959ad35d 100755 --- a/scripts/openmc-track-to-vtk +++ b/scripts/openmc-track-to-vtk @@ -6,14 +6,14 @@ import argparse -import h5py +import openmc import vtk def _parse_args(): # Create argument parser. parser = argparse.ArgumentParser( - description='Convert particle track file to a .pvtp file.') + description='Convert particle track file(s) to a .pvtp file.') parser.add_argument('input', metavar='IN', type=str, nargs='+', help='Input particle track data filename(s).') parser.add_argument('-o', '--out', metavar='OUT', type=str, dest='out', @@ -36,24 +36,25 @@ def main(): # Initialize data arrays and offset. points = vtk.vtkPoints() cells = vtk.vtkCellArray() + point_offset = 0 for fname in args.input: # Write coordinate values to points array. - track = h5py.File(fname) - offsets = track.attrs['offsets'] - for start, end in zip(offsets[:-1], offsets[1:]): - tracks = track['tracks'][start:end] - for arr in tracks: - points.InsertNextPoint(arr['r']) + track_file = openmc.TrackFile(fname) + for track in track_file: + for particle in track.particles: + for state in particle.states: + points.InsertNextPoint(state['r']) - for start, end in zip(offsets[:-1], offsets[1:]): - # Create VTK line and assign points to line. - line = vtk.vtkPolyLine() - line.GetPointIds().SetNumberOfIds(end - start) - for j in range(end - start): - line.GetPointIds().SetId(j, start + j) + # Create VTK line and assign points to line. + n = particle.states.size + line = vtk.vtkPolyLine() + line.GetPointIds().SetNumberOfIds(n) + for i in range(n): + line.GetPointIds().SetId(i, point_offset + i) + point_offset += n - # Add line to cell array - cells.InsertNextCell(line) + # Add line to cell array + cells.InsertNextCell(line) data = vtk.vtkPolyData() data.SetPoints(points) diff --git a/src/simulation.cpp b/src/simulation.cpp index b7f5b58f1..728de8a85 100644 --- a/src/simulation.cpp +++ b/src/simulation.cpp @@ -82,6 +82,11 @@ int openmc_simulation_init() // Allocate source, fission and surface source banks. allocate_banks(); + // Create track file if needed + if (!settings::track_identifiers.empty() || settings::write_all_tracks) { + open_track_file(); + } + // If doing an event-based simulation, intialize the particle buffer // and event queues if (settings::event_based) { @@ -152,6 +157,11 @@ int openmc_simulation_finalize() mat->mat_nuclide_index_.clear(); } + // Close track file if open + if (!settings::track_identifiers.empty() || settings::write_all_tracks) { + close_track_file(); + } + // Increment total number of generations simulation::total_gen += simulation::current_batch * settings::gen_per_batch; diff --git a/src/track_output.cpp b/src/track_output.cpp index 5a7aff834..c75468b1f 100644 --- a/src/track_output.cpp +++ b/src/track_output.cpp @@ -9,6 +9,7 @@ #include "xtensor/xtensor.hpp" #include +#include #include // for size_t #include @@ -19,6 +20,9 @@ namespace openmc { // Global variables //============================================================================== +hid_t track_file; //! HDF5 identifier for track file +hid_t track_dtype; //! HDF5 identifier for track datatype + //============================================================================== // Non-member functions //============================================================================== @@ -34,8 +38,14 @@ void write_particle_track(Particle& p) p.tracks().back().states.push_back(p.get_track_state()); } -hid_t trackstate_type() +void open_track_file() { + // Open file and write filetype/version + std::string filename = fmt::format("{}tracks.h5", settings::path_output); + track_file = file_open(filename, 'w'); + write_attribute(track_file, "filetype", "track"); + write_attribute(track_file, "version", VERSION_TRACK); + // Create compound type for Position hid_t postype = H5Tcreate(H5T_COMPOUND, sizeof(struct Position)); H5Tinsert(postype, "x", HOFFSET(Position, x), H5T_NATIVE_DOUBLE); @@ -43,26 +53,27 @@ hid_t trackstate_type() H5Tinsert(postype, "z", HOFFSET(Position, z), H5T_NATIVE_DOUBLE); // Create compound type for TrackState - hid_t tracktype = H5Tcreate(H5T_COMPOUND, sizeof(struct TrackState)); - H5Tinsert(tracktype, "r", HOFFSET(TrackState, r), postype); - H5Tinsert(tracktype, "u", HOFFSET(TrackState, u), postype); - H5Tinsert(tracktype, "E", HOFFSET(TrackState, E), H5T_NATIVE_DOUBLE); - H5Tinsert(tracktype, "time", HOFFSET(TrackState, time), H5T_NATIVE_DOUBLE); - H5Tinsert(tracktype, "wgt", HOFFSET(TrackState, wgt), H5T_NATIVE_DOUBLE); - H5Tinsert(tracktype, "cell_id", HOFFSET(TrackState, cell_id), H5T_NATIVE_INT); + track_dtype = H5Tcreate(H5T_COMPOUND, sizeof(struct TrackState)); + H5Tinsert(track_dtype, "r", HOFFSET(TrackState, r), postype); + H5Tinsert(track_dtype, "u", HOFFSET(TrackState, u), postype); + H5Tinsert(track_dtype, "E", HOFFSET(TrackState, E), H5T_NATIVE_DOUBLE); + H5Tinsert(track_dtype, "time", HOFFSET(TrackState, time), H5T_NATIVE_DOUBLE); + H5Tinsert(track_dtype, "wgt", HOFFSET(TrackState, wgt), H5T_NATIVE_DOUBLE); H5Tinsert( - tracktype, "material_id", HOFFSET(TrackState, material_id), H5T_NATIVE_INT); - + track_dtype, "cell_id", HOFFSET(TrackState, cell_id), H5T_NATIVE_INT); + H5Tinsert(track_dtype, "material_id", HOFFSET(TrackState, material_id), + H5T_NATIVE_INT); H5Tclose(postype); - return tracktype; +} + +void close_track_file() +{ + H5Tclose(track_dtype); + file_close(track_file); } void finalize_particle_track(Particle& p) { - std::string filename = - fmt::format("{}track_{}_{}_{}.h5", settings::path_output, - simulation::current_batch, simulation::current_gen, p.id()); - // Determine number of coordinates for each particle vector offsets; vector particles; @@ -78,28 +89,25 @@ void finalize_particle_track(Particle& p) #pragma omp critical(FinalizeParticleTrack) { - hid_t file_id = file_open(filename, 'w'); - write_attribute(file_id, "filetype", "track"); - write_attribute(file_id, "version", VERSION_TRACK); - write_attribute(file_id, "n_particles", p.tracks().size()); - write_attribute(file_id, "offsets", offsets); - write_attribute(file_id, "particles", particles); - - // Create HDF5 datatype for TrackState - hid_t track_type = trackstate_type(); + // Create name for dataset + std::string dset_name = fmt::format("track_{}_{}_{}", + simulation::current_batch, simulation::current_gen, p.id()); // Write array of TrackState to file hsize_t dims[] {static_cast(tracks.size())}; hid_t dspace = H5Screate_simple(1, dims, nullptr); - hid_t dset = H5Dcreate(file_id, "tracks", track_type, dspace, H5P_DEFAULT, - H5P_DEFAULT, H5P_DEFAULT); - H5Dwrite(dset, track_type, H5S_ALL, H5S_ALL, H5P_DEFAULT, tracks.data()); + hid_t dset = H5Dcreate(track_file, dset_name.c_str(), track_dtype, dspace, + H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT); + H5Dwrite(dset, track_dtype, H5S_ALL, H5S_ALL, H5P_DEFAULT, tracks.data()); + + // Write attributes + write_attribute(dset, "n_particles", p.tracks().size()); + write_attribute(dset, "offsets", offsets); + write_attribute(dset, "particles", particles); // Free resources H5Dclose(dset); H5Sclose(dspace); - H5Tclose(track_type); - file_close(file_id); } // Clear particle tracks