Introduce WeightWindowsList class that enables export to HDF5 (#3456)

This commit is contained in:
Paul Romano 2025-06-25 13:44:53 -05:00 committed by GitHub
parent a6db05ac8b
commit 01fa8056d1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 337 additions and 210 deletions

View file

@ -37,7 +37,6 @@ Simulation Settings
openmc.read_source_file
openmc.write_source_file
openmc.wwinp_to_wws
Material Specification
----------------------
@ -259,8 +258,16 @@ Variance Reduction
:template: myclass
openmc.WeightWindows
openmc.WeightWindowsList
openmc.WeightWindowGenerator
.. autosummary::
:toctree: generated
:nosignatures:
:template: myfunction.rst
openmc.hdf5_to_wws
openmc.wwinp_to_wws
Coarse Mesh Finite Difference Acceleration

View file

@ -162,7 +162,7 @@ solver, the Python input just needs to load the h5 file::
settings.weight_window_checkpoints = {'collision': True, 'surface': True}
settings.survival_biasing = False
settings.weight_windows = openmc.hdf5_to_wws('weight_windows.h5')
settings.weight_windows = openmc.WeightWindowsList.from_hdf5('weight_windows.h5')
settings.weight_windows_on = True
The :class:`~openmc.WeightWindowGenerator` instance is not needed to load an

View file

@ -16,7 +16,7 @@ from .mesh import _read_meshes, RegularMesh, MeshBase
from .source import SourceBase, MeshSource, IndependentSource
from .utility_funcs import input_path
from .volume import VolumeCalculation
from .weight_windows import WeightWindows, WeightWindowGenerator
from .weight_windows import WeightWindows, WeightWindowGenerator, WeightWindowsList
class RunMode(Enum):
@ -313,7 +313,7 @@ class Settings:
described in :ref:`verbosity`.
volume_calculations : VolumeCalculation or iterable of VolumeCalculation
Stochastic volume calculation specifications
weight_windows : WeightWindows or iterable of WeightWindows
weight_windows : WeightWindowsList
Weight windows to use for variance reduction
.. versionadded:: 0.13
@ -424,7 +424,7 @@ class Settings:
self._max_particles_in_flight = None
self._max_particle_events = None
self._write_initial_source = None
self._weight_windows = cv.CheckedList(WeightWindows, 'weight windows')
self._weight_windows = WeightWindowsList()
self._weight_window_generators = cv.CheckedList(WeightWindowGenerator, 'weight window generators')
self._weight_windows_on = None
self._weight_windows_file = None
@ -1095,14 +1095,14 @@ class Settings:
self._write_initial_source = value
@property
def weight_windows(self) -> list[WeightWindows]:
def weight_windows(self) -> WeightWindowsList:
return self._weight_windows
@weight_windows.setter
def weight_windows(self, value: WeightWindows | Iterable[WeightWindows]):
if not isinstance(value, MutableSequence):
def weight_windows(self, value: WeightWindows | Sequence[WeightWindows]):
if not isinstance(value, Sequence):
value = [value]
self._weight_windows = cv.CheckedList(WeightWindows, 'weight windows', value)
self._weight_windows = WeightWindowsList(value)
@property
def weight_windows_on(self) -> bool:

View file

@ -1,6 +1,8 @@
from __future__ import annotations
from numbers import Real, Integral
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Self
import warnings
import lxml.etree as ET
@ -14,6 +16,7 @@ import openmc.checkvalue as cv
from openmc.checkvalue import PathLike
from ._xml import get_text, clean_indentation
from .mixin import IDManagerMixin
from .utility_funcs import change_directory
class WeightWindows(IDManagerMixin):
@ -90,7 +93,7 @@ class WeightWindows(IDManagerMixin):
survival_ratio : float
Ratio of the survival weight to the lower weight window bound for
rouletting
max_lower_bound_ratio: float
max_lower_bound_ratio : float
Maximum allowed ratio of a particle's weight to the weight window's
lower bound. (Default: 1.0)
max_split : int
@ -114,7 +117,7 @@ class WeightWindows(IDManagerMixin):
upper_bound_ratio: float | None = None,
energy_bounds: Iterable[Real] | None = None,
particle_type: str = 'neutron',
survival_ratio: float = 3,
survival_ratio: float = 3.0,
max_lower_bound_ratio: float | None = None,
max_split: int = 10,
weight_cutoff: float = 1.e-38,
@ -354,7 +357,7 @@ class WeightWindows(IDManagerMixin):
return element
@classmethod
def from_xml_element(cls, elem: ET.Element, meshes: dict[int, MeshBase]) -> WeightWindows:
def from_xml_element(cls, elem: ET.Element, meshes: dict[int, MeshBase]) -> Self:
"""Generate weight window settings from an XML element
Parameters
@ -408,7 +411,7 @@ class WeightWindows(IDManagerMixin):
)
@classmethod
def from_hdf5(cls, group: h5py.Group, meshes: dict[int, MeshBase]) -> WeightWindows:
def from_hdf5(cls, group: h5py.Group, meshes: dict[int, MeshBase]) -> Self:
"""Create weight windows from HDF5 group
Parameters
@ -458,7 +461,7 @@ class WeightWindows(IDManagerMixin):
)
def wwinp_to_wws(path: PathLike) -> list[WeightWindows]:
def wwinp_to_wws(path: PathLike) -> WeightWindowsList:
"""Create WeightWindows instances from a wwinp file
.. versionadded:: 0.13.1
@ -470,190 +473,13 @@ def wwinp_to_wws(path: PathLike) -> list[WeightWindows]:
Returns
-------
list of openmc.WeightWindows
WeightWindowsList
"""
with open(path) as wwinp:
# BLOCK 1
header = wwinp.readline().split(None, 4)
# read file type, time-dependence, number of
# particles, mesh type and problem identifier
_if, iv, ni, nr = [int(x) for x in header[:4]]
# header value checks
if _if != 1:
raise ValueError(f'Found incorrect file type, if: {_if}')
if iv > 1:
# read number of time bins for each particle, 'nt(1...ni)'
nt = np.fromstring(wwinp.readline(), sep=' ', dtype=int)
# raise error if time bins are present for now
raise ValueError('Time-dependent weight windows '
'are not yet supported')
else:
nt = ni * [1]
# read number of energy bins for each particle, 'ne(1...ni)'
ne = np.fromstring(wwinp.readline(), sep=' ', dtype=int)
# read coarse mesh dimensions and lower left corner
mesh_description = np.fromstring(wwinp.readline(), sep=' ')
nfx, nfy, nfz = mesh_description[:3].astype(int)
xyz0 = mesh_description[3:]
# read cylindrical and spherical mesh vectors if present
if nr == 16:
# read number of coarse bins
line_arr = np.fromstring(wwinp.readline(), sep=' ')
ncx, ncy, ncz = line_arr[:3].astype(int)
# read polar vector (x1, y1, z1)
xyz1 = line_arr[3:]
# read azimuthal vector (x2, y2, z2)
line_arr = np.fromstring(wwinp.readline(), sep=' ')
xyz2 = line_arr[:3]
# Get polar and azimuthal axes
polar_axis = xyz1 - xyz0
azimuthal_axis = xyz2 - xyz0
# Check for polar axis other than (0, 0, 1)
norm = np.linalg.norm(polar_axis)
if not np.isclose(polar_axis[2]/norm, 1.0):
raise NotImplementedError('Polar axis not aligned to z-axis not supported')
# Check for azimuthal axis other than (1, 0, 0)
norm = np.linalg.norm(azimuthal_axis)
if not np.isclose(azimuthal_axis[0]/norm, 1.0):
raise NotImplementedError('Azimuthal axis not aligned to x-axis not supported')
# read geometry type
nwg = int(line_arr[-1])
elif nr == 10:
# read rectilinear data:
# number of coarse mesh bins and mesh type
ncx, ncy, ncz, nwg = \
np.fromstring(wwinp.readline(), sep=' ').astype(int)
else:
raise RuntimeError(f'Invalid mesh description (nr) found: {nr}')
# read BLOCK 2 and BLOCK 3 data into a single array
ww_data = np.fromstring(wwinp.read(), sep=' ')
# extract mesh data from the ww_data array
start_idx = 0
# first values in the mesh definition arrays are the first
# coordinate of the grid
end_idx = start_idx + 1 + 3 * ncx
i0, i_vals = ww_data[start_idx], ww_data[start_idx+1:end_idx]
start_idx = end_idx
end_idx = start_idx + 1 + 3 * ncy
j0, j_vals = ww_data[start_idx], ww_data[start_idx+1:end_idx]
start_idx = end_idx
end_idx = start_idx + 1 + 3 * ncz
k0, k_vals = ww_data[start_idx], ww_data[start_idx+1:end_idx]
start_idx = end_idx
# mesh consistency checks
if nr == 16 and nwg == 1 or nr == 10 and nwg != 1:
raise ValueError(f'Mesh description in header ({nr}) '
f'does not match the mesh type ({nwg})')
if nr == 10 and (xyz0 != (i0, j0, k0)).any():
raise ValueError(f'Mesh origin in the header ({xyz0}) '
f' does not match the origin in the mesh '
f' description ({i0, j0, k0})')
# create openmc mesh object
grids = []
mesh_definition = [(i0, i_vals, nfx), (j0, j_vals, nfy), (k0, k_vals, nfz)]
for grid0, grid_vals, n_pnts in mesh_definition:
# file spec checks for the mesh definition
if (grid_vals[2::3] != 1.0).any():
raise ValueError('One or more mesh ratio value, qx, '
'is not equal to one')
s = int(grid_vals[::3].sum())
if s != n_pnts:
raise ValueError(f'Sum of the fine bin entries, {s}, does '
f'not match the number of fine bins, {n_pnts}')
# extend the grid based on the next coarse bin endpoint, px
# and the number of fine bins in the coarse bin, sx
intervals = grid_vals.reshape(-1, 3)
coords = [grid0]
for sx, px, qx in intervals:
coords += np.linspace(coords[-1], px, int(sx + 1)).tolist()[1:]
grids.append(np.array(coords))
if nwg == 1:
mesh = RectilinearMesh()
mesh.x_grid, mesh.y_grid, mesh.z_grid = grids
elif nwg == 2:
mesh = CylindricalMesh(
r_grid=grids[0],
z_grid=grids[1],
phi_grid=grids[2],
origin = xyz0,
)
elif nwg == 3:
mesh = SphericalMesh(
r_grid=grids[0],
theta_grid=grids[1],
phi_grid=grids[2],
origin = xyz0
)
# extract weight window values from array
wws = []
for ne_i, nt_i, particle_type in zip(ne, nt, ('neutron', 'photon')):
# no information to read for this particle if
# either the energy bins or time bins are empty
if ne_i == 0 or nt_i == 0:
continue
if iv > 1:
# time bins are parsed but unused for now
end_idx = start_idx + nt_i
time_bounds = ww_data[start_idx:end_idx]
np.insert(time_bounds, (0,), (0.0,))
start_idx = end_idx
# read energy boundaries
end_idx = start_idx + ne_i
energy_bounds = np.insert(ww_data[start_idx:end_idx], (0,), (0.0,))
# convert from MeV to eV
energy_bounds *= 1e6
start_idx = end_idx
# read weight window values
end_idx = start_idx + (nfx * nfy * nfz) * nt_i * ne_i
# read values and reshape according to ordering
# slowest to fastest: t, e, z, y, x
# reorder with transpose since our ordering is x, y, z, e, t
ww_shape = (nt_i, ne_i, nfz, nfy, nfx)
ww_values = ww_data[start_idx:end_idx].reshape(ww_shape).T
# Only use first time bin since we don't support time dependent weight
# windows yet.
ww_values = ww_values[:, :, :, :, 0]
start_idx = end_idx
# create a weight window object
ww = WeightWindows(id=None,
mesh=mesh,
lower_ww_bounds=ww_values,
upper_bound_ratio=5.0,
energy_bounds=energy_bounds,
particle_type=particle_type)
wws.append(ww)
return wws
warnings.warn(
"This function is deprecated in favor of 'WeightWindowsList.from_wwinp'",
FutureWarning
)
return WeightWindowsList.from_wwinp(path)
class WeightWindowGenerator:
@ -886,7 +712,7 @@ class WeightWindowGenerator:
return element
@classmethod
def from_xml_element(cls, elem: ET.Element, meshes: dict) -> WeightWindowGenerator:
def from_xml_element(cls, elem: ET.Element, meshes: dict) -> Self:
"""
Create a weight window generation object from an XML element
@ -929,8 +755,8 @@ class WeightWindowGenerator:
return wwg
def hdf5_to_wws(path='weight_windows.h5'):
"""Create WeightWindows instances from a weight windows HDF5 file
def hdf5_to_wws(path='weight_windows.h5') -> WeightWindowsList:
"""Create a WeightWindowsList from a weight windows HDF5 file
.. versionadded:: 0.14.0
@ -941,13 +767,284 @@ def hdf5_to_wws(path='weight_windows.h5'):
Returns
-------
list of openmc.WeightWindows
WeightWindowsList
"""
warnings.warn(
"This function is deprecated in favor of 'WeightWindowsList.from_hdf5'",
FutureWarning
)
return WeightWindowsList.from_hdf5(path)
with h5py.File(path) as h5_file:
# read in all of the meshes in the mesh node
meshes = {}
for mesh_group in h5_file['meshes']:
mesh = MeshBase.from_hdf5(h5_file['meshes'][mesh_group])
meshes[mesh.id] = mesh
return [WeightWindows.from_hdf5(ww, meshes) for ww in h5_file['weight_windows'].values()]
class WeightWindowsList(list):
"""A list of WeightWindows objects.
.. versionadded:: 0.15.3
Parameters
----------
iterable : iterable of openmc.WeightWindows
An iterable of WeightWindows objects to initialize the list with
"""
def __init__(self, iterable: Iterable[WeightWindows] = ()):
super().__init__(iterable)
@classmethod
def from_hdf5(cls, path: PathLike = 'weight_windows.h5') -> Self:
"""Create WeightWindowsList from a weight windows HDF5 file.
Parameters
----------
path : PathLike
Path to the weight windows hdf5 file
Returns
-------
WeightWindowsList
A list of WeightWindows objects read from the file
"""
with h5py.File(path) as h5_file:
# read in all of the meshes in the mesh node
meshes = {}
for mesh_group in h5_file['meshes']:
mesh = MeshBase.from_hdf5(h5_file['meshes'][mesh_group])
meshes[mesh.id] = mesh
wws = [
WeightWindows.from_hdf5(ww, meshes)
for ww in h5_file['weight_windows'].values()
]
return cls(wws)
@classmethod
def from_wwinp(cls, path: PathLike) -> Self:
"""Create WeightWindowsList from a wwinp file.
Parameters
----------
path : PathLike
Path to the wwinp file
Returns
-------
WeightWindowsList
A list of WeightWindows objects read from the file
"""
with open(path) as wwinp:
# BLOCK 1
header = wwinp.readline().split(None, 4)
# read file type, time-dependence, number of
# particles, mesh type and problem identifier
_if, iv, ni, nr = [int(x) for x in header[:4]]
# header value checks
if _if != 1:
raise ValueError(f'Found incorrect file type, if: {_if}')
if iv > 1:
# read number of time bins for each particle, 'nt(1...ni)'
nt = np.fromstring(wwinp.readline(), sep=' ', dtype=int)
# raise error if time bins are present for now
raise ValueError('Time-dependent weight windows '
'are not yet supported')
else:
nt = ni * [1]
# read number of energy bins for each particle, 'ne(1...ni)'
ne = np.fromstring(wwinp.readline(), sep=' ', dtype=int)
# read coarse mesh dimensions and lower left corner
mesh_description = np.fromstring(wwinp.readline(), sep=' ')
nfx, nfy, nfz = mesh_description[:3].astype(int)
xyz0 = mesh_description[3:]
# read cylindrical and spherical mesh vectors if present
if nr == 16:
# read number of coarse bins
line_arr = np.fromstring(wwinp.readline(), sep=' ')
ncx, ncy, ncz = line_arr[:3].astype(int)
# read polar vector (x1, y1, z1)
xyz1 = line_arr[3:]
# read azimuthal vector (x2, y2, z2)
line_arr = np.fromstring(wwinp.readline(), sep=' ')
xyz2 = line_arr[:3]
# Get polar and azimuthal axes
polar_axis = xyz1 - xyz0
azimuthal_axis = xyz2 - xyz0
# Check for polar axis other than (0, 0, 1)
norm = np.linalg.norm(polar_axis)
if not np.isclose(polar_axis[2]/norm, 1.0):
raise NotImplementedError('Polar axis not aligned to z-axis not supported')
# Check for azimuthal axis other than (1, 0, 0)
norm = np.linalg.norm(azimuthal_axis)
if not np.isclose(azimuthal_axis[0]/norm, 1.0):
raise NotImplementedError('Azimuthal axis not aligned to x-axis not supported')
# read geometry type
nwg = int(line_arr[-1])
elif nr == 10:
# read rectilinear data:
# number of coarse mesh bins and mesh type
ncx, ncy, ncz, nwg = \
np.fromstring(wwinp.readline(), sep=' ').astype(int)
else:
raise RuntimeError(f'Invalid mesh description (nr) found: {nr}')
# read BLOCK 2 and BLOCK 3 data into a single array
ww_data = np.fromstring(wwinp.read(), sep=' ')
# extract mesh data from the ww_data array
start_idx = 0
# first values in the mesh definition arrays are the first
# coordinate of the grid
end_idx = start_idx + 1 + 3 * ncx
i0, i_vals = ww_data[start_idx], ww_data[start_idx+1:end_idx]
start_idx = end_idx
end_idx = start_idx + 1 + 3 * ncy
j0, j_vals = ww_data[start_idx], ww_data[start_idx+1:end_idx]
start_idx = end_idx
end_idx = start_idx + 1 + 3 * ncz
k0, k_vals = ww_data[start_idx], ww_data[start_idx+1:end_idx]
start_idx = end_idx
# mesh consistency checks
if nr == 16 and nwg == 1 or nr == 10 and nwg != 1:
raise ValueError(f'Mesh description in header ({nr}) '
f'does not match the mesh type ({nwg})')
if nr == 10 and (xyz0 != (i0, j0, k0)).any():
raise ValueError(f'Mesh origin in the header ({xyz0}) '
f' does not match the origin in the mesh '
f' description ({i0, j0, k0})')
# create openmc mesh object
grids = []
mesh_definition = [(i0, i_vals, nfx), (j0, j_vals, nfy), (k0, k_vals, nfz)]
for grid0, grid_vals, n_pnts in mesh_definition:
# file spec checks for the mesh definition
if (grid_vals[2::3] != 1.0).any():
raise ValueError('One or more mesh ratio value, qx, '
'is not equal to one')
s = int(grid_vals[::3].sum())
if s != n_pnts:
raise ValueError(f'Sum of the fine bin entries, {s}, does '
f'not match the number of fine bins, {n_pnts}')
# extend the grid based on the next coarse bin endpoint, px
# and the number of fine bins in the coarse bin, sx
intervals = grid_vals.reshape(-1, 3)
coords = [grid0]
for sx, px, qx in intervals:
coords += np.linspace(coords[-1], px, int(sx + 1)).tolist()[1:]
grids.append(np.array(coords))
if nwg == 1:
mesh = RectilinearMesh()
mesh.x_grid, mesh.y_grid, mesh.z_grid = grids
elif nwg == 2:
mesh = CylindricalMesh(
r_grid=grids[0],
z_grid=grids[1],
phi_grid=grids[2],
origin = xyz0,
)
elif nwg == 3:
mesh = SphericalMesh(
r_grid=grids[0],
theta_grid=grids[1],
phi_grid=grids[2],
origin = xyz0
)
# extract weight window values from array
wws = cls()
for ne_i, nt_i, particle_type in zip(ne, nt, ('neutron', 'photon')):
# no information to read for this particle if
# either the energy bins or time bins are empty
if ne_i == 0 or nt_i == 0:
continue
if iv > 1:
# time bins are parsed but unused for now
end_idx = start_idx + nt_i
time_bounds = ww_data[start_idx:end_idx]
np.insert(time_bounds, (0,), (0.0,))
start_idx = end_idx
# read energy boundaries
end_idx = start_idx + ne_i
energy_bounds = np.insert(ww_data[start_idx:end_idx], (0,), (0.0,))
# convert from MeV to eV
energy_bounds *= 1e6
start_idx = end_idx
# read weight window values
end_idx = start_idx + (nfx * nfy * nfz) * nt_i * ne_i
# read values and reshape according to ordering
# slowest to fastest: t, e, z, y, x
# reorder with transpose since our ordering is x, y, z, e, t
ww_shape = (nt_i, ne_i, nfz, nfy, nfx)
ww_values = ww_data[start_idx:end_idx].reshape(ww_shape).T
# Only use first time bin since we don't support time dependent weight
# windows yet.
ww_values = ww_values[:, :, :, :, 0]
start_idx = end_idx
# create a weight window object
ww = WeightWindows(id=None,
mesh=mesh,
lower_ww_bounds=ww_values,
upper_bound_ratio=5.0,
energy_bounds=energy_bounds,
particle_type=particle_type)
wws.append(ww)
return wws
def export_to_hdf5(self, path: PathLike = 'weight_windows.h5', **init_kwargs):
"""Write weight windows to an HDF5 file.
Parameters
----------
path : PathLike
Path to the file to write weight windows to
**init_kwargs
Keyword arguments passed to :func:`openmc.lib.init`
"""
import openmc.lib
cv.check_type('path', path, PathLike)
# Create a temporary model with the weight windows
model = openmc.Model()
sph = openmc.Sphere(boundary_type='vacuum')
cell = openmc.Cell(region=-sph)
model.geometry = openmc.Geometry([cell])
model.settings.weight_windows = self
model.settings.particles = 100
model.settings.batches = 1
# Get absolute path before moving to temporary directory
path = Path(path).resolve()
with change_directory(tmpdir=True):
# Write the model to an XML file
model.export_to_model_xml()
# Load the model with openmc.lib and then export it to an HDF5 file
with openmc.lib.run_in_memory(**init_kwargs):
openmc.lib.export_weight_windows(path)

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1,23 @@
import openmc
def test_ww_roundtrip(request, run_in_tmpdir):
# Load weight windows from a wwinp file
wwinp_file = request.path.with_name('wwinp_n')
wws = openmc.WeightWindowsList.from_wwinp(wwinp_file)
# Roundtrip them, writing to HDF5 and reading back in
wws.export_to_hdf5('ww.h5')
wws_new = openmc.WeightWindowsList.from_hdf5('ww.h5')
# Check that the new weight windows are the same as the original
assert len(wws) == len(wws_new)
for ww, ww_new in zip(wws, wws_new):
assert ww.particle_type == ww_new.particle_type
assert (ww.lower_ww_bounds == ww_new.lower_ww_bounds).all()
assert (ww.upper_ww_bounds == ww_new.upper_ww_bounds).all()
assert ww.survival_ratio == ww_new.survival_ratio
assert ww.num_energy_bins == ww_new.num_energy_bins
assert ww.max_split == ww_new.max_split
assert ww.weight_cutoff == ww_new.weight_cutoff
assert ww.mesh.id == ww_new.mesh.id