From 2450eef4246cae85c7ee8b895458729b7e6dc97c Mon Sep 17 00:00:00 2001 From: Zoe Prieto <101403129+zoeprieto@users.noreply.github.com> Date: Sat, 5 Oct 2024 03:51:56 -0300 Subject: [PATCH] Introduce ParticleList class for manipulating a list of source particles (#3148) Co-authored-by: Paul Romano --- docs/source/pythonapi/base.rst | 1 + openmc/lib/core.py | 9 +- openmc/source.py | 248 +++++++++++++++--- .../surface_source_write/test.py | 45 ++-- tests/unit_tests/test_source_file.py | 30 ++- 5 files changed, 259 insertions(+), 74 deletions(-) diff --git a/docs/source/pythonapi/base.rst b/docs/source/pythonapi/base.rst index 611d2c216..23df02f2e 100644 --- a/docs/source/pythonapi/base.rst +++ b/docs/source/pythonapi/base.rst @@ -191,6 +191,7 @@ Post-processing :template: myclass.rst openmc.Particle + openmc.ParticleList openmc.ParticleTrack openmc.StatePoint openmc.Summary diff --git a/openmc/lib/core.py b/openmc/lib/core.py index e646a9ae1..8561602e6 100644 --- a/openmc/lib/core.py +++ b/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(): diff --git a/openmc/source.py b/openmc/source.py index e35e62f5f..7c8e6af8d 100644 --- a/openmc/source.py +++ b/openmc/source.py @@ -5,10 +5,12 @@ from enum import IntEnum 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 @@ class ParticleType(IntEnum): 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 @@ class ParticleType(IntEnum): 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', ' list[SourceParticle]: +class ParticleList(list): + """A collection of SourceParticle objects. + + Parameters + ---------- + particles : list of SourceParticle + Particles to collect into the list + + """ + @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 + + """ + 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', ' 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'][...] + filename = Path(filename) + if filename.suffix not in ('.h5', '.mcpl'): + raise ValueError('Source file must have a .h5 or .mcpl extension.') - 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 + if filename.suffix == '.h5': + return ParticleList.from_hdf5(filename) + else: + return ParticleList.from_mcpl(filename) diff --git a/tests/regression_tests/surface_source_write/test.py b/tests/regression_tests/surface_source_write/test.py index fe11d68a6..f144eb82a 100644 --- a/tests/regression_tests/surface_source_write/test.py +++ b/tests/regression_tests/surface_source_write/test.py @@ -608,11 +608,6 @@ def return_surface_source_data(filepath): """Read a surface source file and return a sorted array composed of flatten arrays of source data for each surface source point. - TODO: - - - use read_source_file from source.py instead. Or a dedicated function - to produce sorted list of source points for a given file. - Parameters ---------- filepath : str @@ -629,27 +624,25 @@ def return_surface_source_data(filepath): keys = [] # Read source file - with h5py.File(filepath, "r") as f: - for point in f["source_bank"]: - r = point["r"] - u = point["u"] - e = point["E"] - time = point["time"] - wgt = point["wgt"] - delayed_group = point["delayed_group"] - surf_id = point["surf_id"] - particle = point["particle"] + source = openmc.read_source_file(filepath) - key = ( - f"{r[0]:.10e} {r[1]:.10e} {r[2]:.10e} {u[0]:.10e} {u[1]:.10e} {u[2]:.10e}" - f"{e:.10e} {time:.10e} {wgt:.10e} {delayed_group} {surf_id} {particle}" - ) - - keys.append(key) - - values = [*r, *u, e, time, wgt, delayed_group, surf_id, particle] - assert len(values) == 12 - data.append(values) + for point in source: + r = point.r + u = point.u + e = point.E + time = point.time + wgt = point.wgt + delayed_group = point.delayed_group + surf_id = point.surf_id + particle = point.particle + key = ( + f"{r[0]:.10e} {r[1]:.10e} {r[2]:.10e} {u[0]:.10e} {u[1]:.10e} {u[2]:.10e}" + f"{e:.10e} {time:.10e} {wgt:.10e} {delayed_group} {surf_id} {particle}" + ) + keys.append(key) + values = [*r, *u, e, time, wgt, delayed_group, surf_id, particle] + assert len(values) == 12 + data.append(values) data = np.array(data) keys = np.array(keys) @@ -1129,4 +1122,4 @@ def test_surface_source_cell_dagmc( harness = SurfaceSourceWriteTestHarness( "statepoint.5.h5", model=model, workdir=folder ) - harness.main() \ No newline at end of file + harness.main() diff --git a/tests/unit_tests/test_source_file.py b/tests/unit_tests/test_source_file.py index 1b5549b00..41906c80f 100644 --- a/tests/unit_tests/test_source_file.py +++ b/tests/unit_tests/test_source_file.py @@ -44,11 +44,9 @@ def test_source_file(run_in_tmpdir): assert np.all(arr['delayed_group'] == 0) assert np.all(arr['particle'] == 0) - # Ensure sites read in are consistent - sites = openmc.read_source_file('test_source.h5') + sites = openmc.ParticleList.from_hdf5('test_source.h5') - assert filetype == b'source' xs = np.array([site.r[0] for site in sites]) ys = np.array([site.r[1] for site in sites]) zs = np.array([site.r[2] for site in sites]) @@ -68,6 +66,32 @@ def test_source_file(run_in_tmpdir): p_types = np.array([s.particle for s in sites]) assert np.all(p_types == 0) + # Ensure a ParticleList item is a SourceParticle + site = sites[0] + assert isinstance(site, openmc.SourceParticle) + assert site.E == pytest.approx(n) + + # Ensure site slice read in and exported are consistent + sites_slice = sites[:10] + sites_slice.export_to_hdf5("test_source_slice.h5") + sites_slice = openmc.ParticleList.from_hdf5('test_source_slice.h5') + + assert isinstance(sites_slice, openmc.ParticleList) + assert len(sites_slice) == 10 + E = np.array([s.E for s in sites_slice]) + np.testing.assert_allclose(E, n - np.arange(10)) + + # Ensure site list read in and exported are consistent + df = sites.to_dataframe() + sites_filtered = sites[df[df.E <= 10.0].index.tolist()] + sites_filtered.export_to_hdf5("test_source_filtered.h5") + sites_filtered = openmc.read_source_file('test_source_filtered.h5') + + assert isinstance(sites_filtered, openmc.ParticleList) + assert len(sites_filtered) == 10 + E = np.array([s.E for s in sites_filtered]) + np.testing.assert_allclose(E, np.arange(10, 0, -1)) + def test_wrong_source_attributes(run_in_tmpdir): # Create a source file with animal attributes