Addition of a collision tracking feature (#3417)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

Author: saliba170499

CMakeLists.txt

@@ -338,6 +338,7 @@ list(APPEND libopenmc_SOURCES
   src/cell.cpp
   src/chain.cpp
   src/cmfd_solver.cpp
+  src/collision_track.cpp
   src/cross_sections.cpp
   src/dagmc.cpp
   src/distribution.cpp

docs/source/io_formats/collision_track.rst

@@ -0,0 +1,46 @@
+.. _io_collision_track:
+
+===========================
+Collision Track File Format
+===========================
+
+When collision tracking is enabled with ``mcpl=false`` (the default), OpenMC
+writes binary data to an HDF5 file named ``collision_track.h5``. The same data
+may also be written after each batch when multiple files are requested
+(``collision_track.N.h5``) or when the run is performed in parallel. The file
+contains the information needed to reconstruct each recorded collision.
+
+The current revision of the collision track file format is 1.0.
+
+**/**
+
+:Attributes:
+            - **filetype** (*char[]*) -- String indicating the type of file.
+              For collision-track files the value is ``"collision_track"``.
+
+:Datasets:
+
+          - **collision_track_bank** (Compound type) -- Collision information
+            for each stored event. Each entry in the dataset corresponds to one
+            collision and contains the following fields:
+
+              - ``r`` (*double[3]*) -- Position of the collision in [cm].
+              - ``u`` (*double[3]*) -- Direction unit vector immediately after the collision.
+              - ``E`` (*double*) -- Incident particle energy before the collision in [eV].
+              - ``dE`` (*double*) -- Energy loss over the collision (:math:`E_\text{before} - E_\text{after}`) in [eV].
+              - ``time`` (*double*) -- Time of the collision in [s].
+              - ``wgt`` (*double*) -- Particle weight at the collision.
+              - ``event_mt`` (*int*) -- ENDF MT number identifying the reaction.
+              - ``delayed_group`` (*int*) -- Delayed neutron group index (non-zero for delayed events).
+              - ``cell_id`` (*int*) -- ID of the cell in which the collision occurred.
+              - ``nuclide_id`` (*int*) -- ZA identifier of the nuclide (ZZZAAAM format).
+              - ``material_id`` (*int*) -- ID of the material containing the collision site.
+              - ``universe_id`` (*int*) -- ID of the universe containing the collision site.
+              - ``n_collision`` (*int*) -- Collision counter for the particle history.
+              - ``particle`` (*int*) -- Particle type (0=neutron, 1=photon, 2=electron, 3=positron).
+              - ``parent_id`` (*int64*) -- Unique ID of the parent particle.
+              - ``progeny_id`` (*int64*) -- Progeny ID of the particle.
+
+In an MPI run, OpenMC writes the combined dataset by gathering collision-track
+entries from all ranks before flushing them to disk, so the final file appears
+as though it were produced serially.

docs/source/io_formats/index.rst

@@ -44,6 +44,7 @@ Output Files
 
    statepoint
    source
+   collision_track
    summary
    properties
    depletion_results

docs/source/io_formats/settings.rst

@@ -20,6 +20,85 @@ source neutrons.
 
   *Default*: None
 
+-----------------------------
+``<collision_track>`` Element
+-----------------------------
+
+The ``<collision_track>`` element indicates to track information about particle
+collisions based on a set of criteria and store these events in a file named
+``collision_track.h5``. This file records details such as the position of the
+interaction, direction of the incoming particle, incident energy and deposited
+energy, weight, time of the interaction, and the delayed neutron group (0 for
+prompt neutrons). Additional information such as the cell ID, material ID,
+universe ID, nuclide ZAID, particle type, and event MT number are also stored.
+Users can specify one or more criterion to filter collisions. If no criteria are
+specified, it defaults to tracking all collisions across the model.
+
+.. warning::
+    Storing all collisions can be very memory intensive. For more targeted
+    tracking, users can employ a variety of parameters such as ``cell_ids``,
+    ``reactions``, ``universe_ids``, ``material_ids``, ``nuclides``, and
+    ``deposited_E_threshold`` to refine the selection of particle interactions
+    to be banked.
+
+This element can contain one or more of the following attributes or
+sub-elements:
+
+  :max_collisions:
+    An integer indicating the maximum number of collisions to be banked per file.
+
+    *Default*: 1000
+
+  :max_collision_track_files:
+    An integer indicating the number of collision_track files to be used.
+
+    *Default*: 1
+
+  :mcpl:
+    An optional boolean to enable MCPL_-format instead of the native HDF5-based
+    format. If activated, the output file name and type is changed to
+    ``collision_track.mcpl``.
+
+    *Default*: false
+
+    .. _MCPL: https://mctools.github.io/mcpl/mcpl.pdf
+
+  :cell_ids:
+    A list of integers representing cell IDs to define specific cells in which
+    collisions are to be banked.
+
+    *Default*: None
+
+  :universe_ids:
+    A list of integers representing the universe IDs to define specific
+    universes in which collisions are to be banked.
+
+    *Default*: None
+
+  :material_ids:
+    A list of integers representing the material IDs to define specific
+    materials in which collisions are to be banked.
+
+    *Default*: None
+
+  :nuclides:
+    A list of strings representing the nuclide, to define specific
+    define specific target nuclide collisions to be banked.
+
+    *Default*: None
+
+  :reactions:
+    A list of integers representing the ENDF-6 format MT numbers or strings
+    (e.g. (n,fission)) to define specific reaction types to be banked.
+
+    *Default*: None
+
+  :deposited_E_threshold:
+    A float defining the minimum deposited energy per collision (in eV) to
+    trigger banking.
+
+    *Default*: 0.0
+
 ----------------------------------
 ``<confidence_intervals>`` Element
 ----------------------------------

docs/source/pythonapi/base.rst

@@ -216,6 +216,9 @@ Post-processing
    :nosignatures:
    :template: myfunction.rst
 
+   openmc.read_collision_track_file
+   openmc.read_collision_track_hdf5
+   openmc.read_collision_track_mcpl
    openmc.voxel_to_vtk
 
 The following classes and functions are used for functional expansion reconstruction.

docs/source/usersguide/settings.rst

@@ -756,6 +756,62 @@ instance, whereas the :meth:`openmc.Track.filter` method returns a new
             track_files = [f"tracks_p{rank}.h5" for rank in range(32)]
             openmc.Tracks.combine(track_files, "tracks.h5")
 
+Collision Track File
+---------------------
+
+OpenMC can generate a collision track file that contains detailed collision
+information (position, direction, energy, deposited energy, time, weight, cell
+ID, material ID, universe ID, nuclide ZAID, particle type, particle delayed
+group and particle ID) for each particle collision depending on user-defined
+parameters. To invoke this feature, set the
+:attr:`~openmc.Settings.collision_track` attribute as shown in this example::
+
+  settings.collision_track = {
+      "max_collisions": 300,
+      "reactions": ["(n,fission)", "(n,2n)"],
+      "material_ids": [1,2],
+      "nuclides": ["U238", "O16"],
+      "cell_ids": [5, 12]
+  }
+
+In this example, collision track information is written to the
+collision_track.h5 file at the end of the simulation. The file contains
+300 recorded collisions that occurred in materials with IDs 1 or 2, involving
+fission or (n,2n) reactions on the nuclides U-238 or O-16, within cells
+with IDs 5 and 12.
+The file can be read using :func:`openmc.read_collision_track_file`.
+The example below shows how to extract the data from the collision_track
+feature and displays the fields stored in the file:
+
+>>> data = openmc.read_collision_track_file('collision_track.h5')
+>>> data.dtype
+    dtype([('r', [('x', '<f8'), ('y', '<f8'), ('z', '<f8')]),
+    ('u', [('x', '<f8'), ('y', '<f8'), ('z', '<f8')]), ('E', '<f8'),
+    ('dE', '<f8'), ('time', '<f8'), ('wgt', '<f8'), ('event_mt', '<i4'),
+    ('delayed_group', '<i4'), ('cell_id', '<i4'), ('nuclide_id', '<i4'),
+    ('material_id', '<i4'), ('universe_id', '<i4'), ('n_collision', '<i4'),
+    ('particle', '<i4'), ('parent_id', '<i8'), ('progeny_id', '<i8')])
+
+
+The full list of fields is as follows:
+
+  :r: Position (each direction in [cm])
+  :u: Direction
+  :E: Energy in [eV]
+  :dE: Energy deposited during collision in [eV]
+  :time: Time in [s]
+  :wgt: Weight of the particle
+  :event_mt: Reaction MT number
+  :delayed_group: Delayed group of the particle
+  :cell_id: Cell ID
+  :nuclide_id: Nuclide ID (10000×Z + 10×A + M)
+  :material_id: Material ID
+  :universe_id: Universe ID
+  :n_collision: Number of collision suffered by the particle
+  :particle: Particle type
+  :parent_id: Source particle ID
+  :progeny_id: Progeny ID
+
 -----------------------
 Restarting a Simulation
 -----------------------

include/openmc/bank.h

@@ -20,6 +20,8 @@ extern vector<SourceSite> source_bank;
 
 extern SharedArray<SourceSite> surf_source_bank;
 
+extern SharedArray<CollisionTrackSite> collision_track_bank;
+
 extern SharedArray<SourceSite> fission_bank;
 
 extern vector<vector<int>> ifp_source_delayed_group_bank;

include/openmc/bank_io.h

@@ -0,0 +1,103 @@
+#ifndef OPENMC_BANK_IO_H
+#define OPENMC_BANK_IO_H
+
+#include "hdf5.h"
+
+#include "openmc/message_passing.h"
+#include "openmc/span.h"
+#include "openmc/vector.h"
+
+#include <algorithm>
+
+#ifdef OPENMC_MPI
+#include <mpi.h>
+#endif
+
+namespace openmc {
+
+template<typename SiteType>
+void write_bank_dataset(const char* dataset_name, hid_t group_id,
+  span<SiteType> bank, const vector<int64_t>& bank_index, hid_t banktype
+#ifdef OPENMC_MPI
+  ,
+  MPI_Datatype mpi_dtype
+#endif
+)
+{
+  int64_t dims_size = bank_index.back();
+  int64_t count_size = bank_index[mpi::rank + 1] - bank_index[mpi::rank];
+
+#ifdef PHDF5
+  hsize_t dims[] {static_cast<hsize_t>(dims_size)};
+  hid_t dspace = H5Screate_simple(1, dims, nullptr);
+  hid_t dset = H5Dcreate(group_id, dataset_name, banktype, dspace, H5P_DEFAULT,
+    H5P_DEFAULT, H5P_DEFAULT);
+
+  hsize_t count[] {static_cast<hsize_t>(count_size)};
+  hid_t memspace = H5Screate_simple(1, count, nullptr);
+
+  hsize_t start[] {static_cast<hsize_t>(bank_index[mpi::rank])};
+  H5Sselect_hyperslab(dspace, H5S_SELECT_SET, start, nullptr, count, nullptr);
+
+  hid_t plist = H5Pcreate(H5P_DATASET_XFER);
+  H5Pset_dxpl_mpio(plist, H5FD_MPIO_COLLECTIVE);
+
+  H5Dwrite(dset, banktype, memspace, dspace, plist, bank.data());
+
+  H5Sclose(dspace);
+  H5Sclose(memspace);
+  H5Dclose(dset);
+  H5Pclose(plist);
+#else
+  if (mpi::master) {
+    hsize_t dims[] {static_cast<hsize_t>(dims_size)};
+    hid_t dspace = H5Screate_simple(1, dims, nullptr);
+    hid_t dset = H5Dcreate(group_id, dataset_name, banktype, dspace,
+      H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
+
+#ifdef OPENMC_MPI
+    vector<SiteType> temp_bank {bank.begin(), bank.end()};
+#endif
+
+    for (int i = 0; i < mpi::n_procs; ++i) {
+      hsize_t count[] {static_cast<hsize_t>(bank_index[i + 1] - bank_index[i])};
+      hid_t memspace = H5Screate_simple(1, count, nullptr);
+
+#ifdef OPENMC_MPI
+      if (i > 0) {
+        MPI_Recv(bank.data(), count[0], mpi_dtype, i, i, mpi::intracomm,
+          MPI_STATUS_IGNORE);
+      }
+#endif
+
+      hid_t dspace_rank = H5Dget_space(dset);
+      hsize_t start[] {static_cast<hsize_t>(bank_index[i])};
+      H5Sselect_hyperslab(
+        dspace_rank, H5S_SELECT_SET, start, nullptr, count, nullptr);
+
+      H5Dwrite(dset, banktype, memspace, dspace_rank, H5P_DEFAULT, bank.data());
+
+      H5Sclose(memspace);
+      H5Sclose(dspace_rank);
+    }
+
+    H5Dclose(dset);
+
+#ifdef OPENMC_MPI
+    std::copy(temp_bank.begin(), temp_bank.end(), bank.begin());
+#endif
+  }
+#ifdef OPENMC_MPI
+  else {
+    if (!bank.empty()) {
+      MPI_Send(
+        bank.data(), bank.size(), mpi_dtype, 0, mpi::rank, mpi::intracomm);
+    }
+  }
+#endif
+#endif
+}
+
+} // namespace openmc
+
+#endif // OPENMC_BANK_IO_H

include/openmc/collision_track.h

@@ -0,0 +1,23 @@
+#ifndef OPENMC_COLLISION_TRACK_H
+#define OPENMC_COLLISION_TRACK_H
+
+#include <string>
+
+namespace openmc {
+
+class Particle;
+
+//! Reserve space in the collision track bank according to user settings.
+void collision_track_reserve_bank();
+
+//! Write collision track data to disk when the bank is full or the batch ends.
+void collision_track_flush_bank();
+
+//! Record the current particle as a collision-track entry when applicable.
+//!
+//! \param particle Particle whose collision should be recorded if eligible
+void collision_track_record(Particle& particle);
+
+} // namespace openmc
+
+#endif // OPENMC_COLLISION_TRACK_H

include/openmc/constants.h

@@ -34,6 +34,7 @@ constexpr array<int, 2> VERSION_VOXEL {2, 0};
 constexpr array<int, 2> VERSION_MGXS_LIBRARY {1, 0};
 constexpr array<int, 2> VERSION_PROPERTIES {1, 1};
 constexpr array<int, 2> VERSION_WEIGHT_WINDOWS {1, 0};
+constexpr array<int, 2> VERSION_COLLISION_TRACK {1, 0};
 
 // ============================================================================
 // ADJUSTABLE PARAMETERS