Write all tracks to a single file

This commit is contained in:
Paul Romano 2022-05-17 07:33:55 -04:00
parent 2e9560a72c
commit 365aa0b0eb
6 changed files with 130 additions and 62 deletions

View file

@ -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.

View file

@ -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

View file

@ -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'<TrackFile: {len(self.tracks)} particles>'
return f'<Track {self.identifier}: {len(self.particles)} particles>'
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))

View file

@ -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)

View file

@ -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;

View file

@ -9,6 +9,7 @@
#include "xtensor/xtensor.hpp"
#include <fmt/core.h>
#include <hdf5.h>
#include <cstddef> // for size_t
#include <string>
@ -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<int> offsets;
vector<int> 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<hsize_t>(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