diff --git a/docs/source/pythonapi/base.rst b/docs/source/pythonapi/base.rst index b471d12a8..fdd89b844 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.TrackFile The following classes and functions are used for functional expansion reconstruction. diff --git a/openmc/__init__.py b/openmc/__init__.py index 0a11fbf53..91af25a5f 100644 --- a/openmc/__init__.py +++ b/openmc/__init__.py @@ -30,6 +30,7 @@ from openmc.mixin import * from openmc.plotter import * from openmc.search import * from openmc.polynomial import * +from openmc.trackfile import * from . import examples # Import a few names from the model module diff --git a/openmc/source.py b/openmc/source.py index 1abe01f65..572333e37 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -298,6 +298,10 @@ class SourceParticle: self.surf_id = surf_id self.particle = particle + def __repr__(self): + name = self.particle.name.lower() + return f'' + def to_tuple(self): """Return source particle attributes as a tuple diff --git a/openmc/trackfile.py b/openmc/trackfile.py new file mode 100644 index 000000000..ab2b81a54 --- /dev/null +++ b/openmc/trackfile.py @@ -0,0 +1,72 @@ +from collections import namedtuple + +import h5py + +from .source import SourceParticle, ParticleType + + +TrackHistory = namedtuple('TrackHistory', ['particle', 'states']) + + +class TrackFile: + """Particle track output + + Parameters + ---------- + filepath : str or pathlib.Path + Path of file to load + + Attributes + ---------- + 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'] + + # 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 + + def __repr__(self): + 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: + r = states['r'] + ax.plot3D(r['x'], r['y'], r['z']) + ax.set_xlabel('x [cm]') + ax.set_ylabel('y [cm]') + ax.set_zlabel('z [cm]') + plt.show() + + @property + def sources(self): + sources = [] + for track_history in self.tracks: + particle_type = ParticleType(track_history.particle) + state = track_history.states[0] + sources.append( + SourceParticle( + r=state['r'], u=state['u'], E=state['E'], + time=state['time'], wgt=state['wgt'], + particle=particle_type + ) + ) + return sources