Introduce ParticleList class for manipulating a list of source partic…

…les (#3148)

Co-authored-by: Paul Romano <[email protected]>

docs/source/pythonapi/base.rst

@@ -191,6 +191,7 @@ Post-processing
    :template: myclass.rst
 
    openmc.Particle
+   openmc.ParticleList
    openmc.ParticleTrack
    openmc.StatePoint
    openmc.Summary

openmc/lib/core.py

@@ -477,7 +477,7 @@ def run(output=True):
 def sample_external_source(
         n_samples: int = 1000,
         prn_seed: int | None = None
-) -> list[openmc.SourceParticle]:
+) -> openmc.ParticleList:
     """Sample external source and return source particles.
 
     .. versionadded:: 0.13.1
@@ -492,7 +492,7 @@ def sample_external_source(
 
     Returns
     -------
-    list of openmc.SourceParticle
+    openmc.ParticleList
         List of sampled source particles
 
     """
@@ -506,14 +506,13 @@ def sample_external_source(
     _dll.openmc_sample_external_source(c_size_t(n_samples), c_uint64(prn_seed), sites_array)
 
     # Convert to list of SourceParticle and return
-    return [
-        openmc.SourceParticle(
+    return openmc.ParticleList([openmc.SourceParticle(
             r=site.r, u=site.u, E=site.E, time=site.time, wgt=site.wgt,
             delayed_group=site.delayed_group, surf_id=site.surf_id,
             particle=openmc.ParticleType(site.particle)
         )
         for site in sites_array
-    ]
+    ])
 
 
 def simulation_init():

openmc/source.py

@@ -5,10 +5,12 @@
 from numbers import Real
 import warnings
 from typing import Any
+from pathlib import Path
 
 import lxml.etree as ET
 import numpy as np
 import h5py
+import pandas as pd
 
 import openmc
 import openmc.checkvalue as cv
@@ -917,6 +919,34 @@ def from_string(cls, value: str):
         except KeyError:
             raise ValueError(f"Invalid string for creation of {cls.__name__}: {value}")
 
+    @classmethod
+    def from_pdg_number(cls, pdg_number: int) -> ParticleType:
+        """Constructs a ParticleType instance from a PDG number.
+
+        The Particle Data Group at LBNL publishes a Monte Carlo particle
+        numbering scheme as part of the `Review of Particle Physics
+        <10.1103/PhysRevD.110.030001>`_. This method maps PDG numbers to the
+        corresponding :class:`ParticleType`.
+
+        Parameters
+        ----------
+        pdg_number : int
+            The PDG number of the particle type.
+
+        Returns
+        -------
+        The corresponding ParticleType instance.
+        """
+        try:
+            return {
+                2112: ParticleType.NEUTRON,
+                22: ParticleType.PHOTON,
+                11: ParticleType.ELECTRON,
+                -11: ParticleType.POSITRON,
+            }[pdg_number]
+        except KeyError:
+            raise ValueError(f"Unrecognized PDG number: {pdg_number}")
+
     def __repr__(self) -> str:
         """
         Returns a string representation of the ParticleType instance.
@@ -930,11 +960,6 @@ def __repr__(self) -> str:
     def __str__(self) -> str:
         return self.__repr__()
 
-    # needed for <= 3.7, IntEnum will use the mixed-in type's `__format__` method otherwise
-    # this forces it to default to the standard object format, relying on __str__ under the hood
-    def __format__(self, spec):
-        return object.__format__(self, spec)
-
 
 class SourceParticle:
     """Source particle
@@ -1020,31 +1045,179 @@ def write_source_file(
     openmc.SourceParticle
 
     """
-    # Create compound datatype for source particles
-    pos_dtype = np.dtype([('x', '<f8'), ('y', '<f8'), ('z', '<f8')])
-    source_dtype = np.dtype([
-        ('r', pos_dtype),
-        ('u', pos_dtype),
-        ('E', '<f8'),
-        ('time', '<f8'),
-        ('wgt', '<f8'),
-        ('delayed_group', '<i4'),
-        ('surf_id', '<i4'),
-        ('particle', '<i4'),
-    ])
-
-    # Create array of source particles
     cv.check_iterable_type("source particles", source_particles, SourceParticle)
-    arr = np.array([s.to_tuple() for s in source_particles], dtype=source_dtype)
+    pl = ParticleList(source_particles)
+    pl.export_to_hdf5(filename, **kwargs)
+
+
+class ParticleList(list):
+    """A collection of SourceParticle objects.
+
+    Parameters
+    ----------
+    particles : list of SourceParticle
+        Particles to collect into the list
 
-    # Write array to file
-    kwargs.setdefault('mode', 'w')
-    with h5py.File(filename, **kwargs) as fh:
-        fh.attrs['filetype'] = np.bytes_("source")
-        fh.create_dataset('source_bank', data=arr, dtype=source_dtype)
+    """
+    @classmethod
+    def from_hdf5(cls, filename: PathLike) -> ParticleList:
+        """Create particle list from an HDF5 file.
 
+        Parameters
+        ----------
+        filename : path-like
+            Path to source file to read.
+
+        Returns
+        -------
+        ParticleList instance
 
-def read_source_file(filename: PathLike) -> list[SourceParticle]:
+        """
+        with h5py.File(filename, 'r') as fh:
+            filetype = fh.attrs['filetype']
+            arr = fh['source_bank'][...]
+
+        if filetype != b'source':
+            raise ValueError(f'File {filename} is not a source file')
+
+        source_particles = [
+            SourceParticle(*params, ParticleType(particle))
+            for *params, particle in arr
+        ]
+        return cls(source_particles)
+
+    @classmethod
+    def from_mcpl(cls, filename: PathLike) -> ParticleList:
+        """Create particle list from an MCPL file.
+
+        Parameters
+        ----------
+        filename : path-like
+            Path to MCPL file to read.
+
+        Returns
+        -------
+        ParticleList instance
+
+        """
+        import mcpl
+        # Process .mcpl file
+        particles = []
+        with mcpl.MCPLFile(filename) as f:
+            for particle in f.particles:
+                # Determine particle type based on the PDG number
+                try:
+                    particle_type = ParticleType.from_pdg_number(particle.pdgcode)
+                except ValueError:
+                    particle_type = "UNKNOWN"
+
+                # Create a source particle instance. Note that MCPL stores
+                # energy in MeV and time in ms.
+                source_particle = SourceParticle(
+                    r=tuple(particle.position),
+                    u=tuple(particle.direction),
+                    E=1.0e6*particle.ekin,
+                    time=1.0e-3*particle.time,
+                    wgt=particle.weight,
+                    particle=particle_type
+                )
+                particles.append(source_particle)
+
+        return cls(particles)
+
+    def __getitem__(self, index):
+        """
+        Return a new ParticleList object containing the particle(s)
+        at the specified index or slice.
+
+        Parameters
+        ----------
+        index : int, slice or list
+            The index, slice or list to select from the list of particles
+
+        Returns
+        -------
+        openmc.ParticleList or openmc.SourceParticle
+            A new object with the selected particle(s)
+        """
+        if isinstance(index, int):
+            # If it's a single integer, return the corresponding particle
+            return super().__getitem__(index)
+        elif isinstance(index, slice):
+            # If it's a slice, return a new ParticleList object with the
+            # sliced particles
+            return ParticleList(super().__getitem__(index))
+        elif isinstance(index, list):
+            # If it's a list of integers, return a new ParticleList object with
+            # the selected particles. Note that Python 3.10 gets confused if you
+            # use super() here, so we call list.__getitem__ directly.
+            return ParticleList([list.__getitem__(self, i) for i in index])
+        else:
+            raise TypeError(f"Invalid index type: {type(index)}. Must be int, "
+                            "slice, or list of int.")
+
+    def to_dataframe(self) -> pd.DataFrame:
+        """A dataframe representing the source particles
+
+        Returns
+        -------
+        pandas.DataFrame
+            DataFrame containing the source particles attributes.
+        """
+        # Extract the attributes of the source particles into a list of tuples
+        data = [(sp.r[0], sp.r[1], sp.r[2], sp.u[0], sp.u[1], sp.u[2],
+                 sp.E, sp.time, sp.wgt, sp.delayed_group, sp.surf_id,
+                 sp.particle.name.lower()) for sp in self]
+
+        # Define the column names for the DataFrame
+        columns = ['x', 'y', 'z', 'u_x', 'u_y', 'u_z', 'E', 'time', 'wgt',
+                   'delayed_group', 'surf_id', 'particle']
+
+        # Create the pandas DataFrame from the data
+        return pd.DataFrame(data, columns=columns)
+
+    def export_to_hdf5(self, filename: PathLike, **kwargs):
+        """Export particle list to an HDF5 file.
+
+        This method write out an .h5 file that can be used as a source file in
+        conjunction with the :class:`openmc.FileSource` class.
+
+        Parameters
+        ----------
+        filename : path-like
+            Path to source file to write
+        **kwargs
+            Keyword arguments to pass to :class:`h5py.File`
+
+        See Also
+        --------
+        openmc.FileSource
+
+        """
+        # Create compound datatype for source particles
+        pos_dtype = np.dtype([('x', '<f8'), ('y', '<f8'), ('z', '<f8')])
+        source_dtype = np.dtype([
+            ('r', pos_dtype),
+            ('u', pos_dtype),
+            ('E', '<f8'),
+            ('time', '<f8'),
+            ('wgt', '<f8'),
+            ('delayed_group', '<i4'),
+            ('surf_id', '<i4'),
+            ('particle', '<i4'),
+        ])
+
+        # Create array of source particles
+        arr = np.array([s.to_tuple() for s in self], dtype=source_dtype)
+
+        # Write array to file
+        kwargs.setdefault('mode', 'w')
+        with h5py.File(filename, **kwargs) as fh:
+            fh.attrs['filetype'] = np.bytes_("source")
+            fh.create_dataset('source_bank', data=arr, dtype=source_dtype)
+
+
+def read_source_file(filename: PathLike) -> ParticleList:
     """Read a source file and return a list of source particles.
 
     .. versionadded:: 0.15.0
@@ -1056,23 +1229,18 @@ def read_source_file(filename: PathLike) -> list[SourceParticle]:
 
     Returns
     -------
-    list of SourceParticle
-        Source particles read from file
+    openmc.ParticleList
 
     See Also
     --------
     openmc.SourceParticle
 
     """
-    with h5py.File(filename, 'r') as fh:
-        filetype = fh.attrs['filetype']
-        arr = fh['source_bank'][...]
-
-    if filetype != b'source':
-        raise ValueError(f'File {filename} is not a source file')
-
-    source_particles = []
-    for *params, particle in arr:
-        source_particles.append(SourceParticle(*params, ParticleType(particle)))
-
-    return source_particles
+    filename = Path(filename)
+    if filename.suffix not in ('.h5', '.mcpl'):
+        raise ValueError('Source file must have a .h5 or .mcpl extension.')
+
+    if filename.suffix == '.h5':
+        return ParticleList.from_hdf5(filename)
+    else:
+        return ParticleList.from_mcpl(filename)