mirror of
https://github.com/openmc-dev/openmc.git
synced 2026-07-28 22:26:08 -04:00
Introduce ParticleList class for manipulating a list of source particles (#3148)
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
This commit is contained in:
Zoe Prieto
GitHub
parent
9070b8b220
commit
2450eef424
5 changed files with 259 additions and 74 deletions
|
|
@ -191,6 +191,7 @@ Post-processing
|
|||
:template: myclass.rst
|
||||
|
||||
openmc.Particle
|
||||
openmc.ParticleList
|
||||
openmc.ParticleTrack
|
||||
openmc.StatePoint
|
||||
openmc.Summary
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
248
openmc/source.py
248
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', '<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)
|
||||
|
||||
# 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)
|
||||
pl = ParticleList(source_particles)
|
||||
pl.export_to_hdf5(filename, **kwargs)
|
||||
|
||||
|
||||
def read_source_file(filename: PathLike) -> 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', '<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'][...]
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
harness.main()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue