From aa909157a009aed7987fc477a01d8ce1e01372f2 Mon Sep 17 00:00:00 2001 From: Perry Date: Fri, 22 May 2026 12:15:17 -0700 Subject: [PATCH 1/2] Speed up check_iterable_type for numpy float/complex arrays Float/complex ndarrays are dtype-validated, so the per-element isinstance() scan is redundant. Also construct upper_ww_bounds in WeightWindows.__init__ as an ndarray multiplication (not a list comprehension) so the upper-bounds setter benefits too. ~11x speedup on 172M-element wwinp inputs (397 s -> 35 s). --- openmc/checkvalue.py | 11 +++++++++++ openmc/weight_windows.py | 4 +--- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/openmc/checkvalue.py b/openmc/checkvalue.py index 5ff2cf9ac5a..e368494de3e 100644 --- a/openmc/checkvalue.py +++ b/openmc/checkvalue.py @@ -80,7 +80,18 @@ def check_iterable_type(name, value, expected_type, min_depth=1, max_depth=1): max_depth : int The maximum number of layers of nested iterables there should be before reaching the ultimately contained items + + Notes + ----- + numpy float/complex ndarrays whose number of dimensions falls within + [``min_depth``, ``max_depth``], the dtype is trusted to guarantee element + type and the per-element scan is skipped, which allows faster processing. """ + # Fast path: float/complex ndarrays of correct depth are dtype-validated. + if (isinstance(value, np.ndarray) and value.dtype.kind in 'fc' + and min_depth <= value.ndim <= max_depth): + return + # Initialize the tree at the very first item. tree = [value] index = [0] diff --git a/openmc/weight_windows.py b/openmc/weight_windows.py index 63af2596efc..99bcc12d019 100644 --- a/openmc/weight_windows.py +++ b/openmc/weight_windows.py @@ -140,9 +140,7 @@ def __init__( "upper_bound_ratio must be present.") if upper_bound_ratio: - self.upper_ww_bounds = [ - lb * upper_bound_ratio for lb in self.lower_ww_bounds - ] + self.upper_ww_bounds = self.lower_ww_bounds * upper_bound_ratio if upper_ww_bounds is not None: self.upper_ww_bounds = upper_ww_bounds From d3124ee5492456b3a7c81665ab2766788794f356 Mon Sep 17 00:00:00 2001 From: Perry Date: Fri, 22 May 2026 12:15:36 -0700 Subject: [PATCH 2/2] Bypass XML round-trip in WeightWindowsList.export_to_hdf5 The XML serialization raised MemoryError on bound arrays >~200M elements -- lxml's intermediate ASCII allocation fails before the text node can be built. Write HDF5 directly via h5py, mirroring the C++ WeightWindows::to_hdf5 writer. Critical details for C++ compatibility: - Bounds are 2D (ne, n_voxels) on disk (4D would segfault the C++ tensor::Tensor reader). - max_lower_bound_ratio is written unconditionally (default 1.0). - Root attrs filetype and version are required by openmc_weight_windows_import. Mesh writing is handled by a private _write_mesh_group helper in weight_windows.py that dispatches by mesh type, matching the reference implementation. UnstructuredMesh raises NotImplementedError (wwinp cannot produce one). --- openmc/weight_windows.py | 121 +++++++++++++++--- .../unit_tests/weightwindows/test_ww_list.py | 23 ++++ 2 files changed, 126 insertions(+), 18 deletions(-) diff --git a/openmc/weight_windows.py b/openmc/weight_windows.py index 99bcc12d019..7c24a4b5920 100644 --- a/openmc/weight_windows.py +++ b/openmc/weight_windows.py @@ -10,7 +10,10 @@ import h5py import openmc -from openmc.mesh import MeshBase, RectilinearMesh, CylindricalMesh, SphericalMesh, UnstructuredMesh +from openmc.mesh import ( + MeshBase, RegularMesh, RectilinearMesh, CylindricalMesh, SphericalMesh, + UnstructuredMesh, +) from openmc.tallies import Tallies import openmc.checkvalue as cv from openmc.checkvalue import PathLike @@ -819,6 +822,51 @@ def hdf5_to_wws(path='weight_windows.h5') -> WeightWindowsList: return WeightWindowsList.from_hdf5(path) +def _write_mesh_group(meshes_grp: h5py.Group, mesh: MeshBase): + """Write *mesh* into *meshes_grp* as a ``mesh `` subgroup. + + Matches the layout produced by ``Mesh::to_hdf5`` in src/mesh.cpp so the + resulting file is readable by ``openmc_weight_windows_import``. + """ + grp = meshes_grp.create_group(f'mesh {mesh.id}') + grp.attrs['id'] = np.int32(mesh.id) + grp.create_dataset('name', data=np.bytes_(mesh.name or '')) + + if isinstance(mesh, RegularMesh): + grp.create_dataset('type', data=np.bytes_('regular')) + grp.create_dataset('dimension', + data=np.asarray(mesh.dimension, dtype=np.int32)) + grp.create_dataset('lower_left', + data=np.asarray(mesh.lower_left, dtype=float)) + grp.create_dataset('upper_right', + data=np.asarray(mesh.upper_right, dtype=float)) + grp.create_dataset('width', data=np.asarray(mesh.width, dtype=float)) + elif isinstance(mesh, RectilinearMesh): + grp.create_dataset('type', data=np.bytes_('rectilinear')) + grp.create_dataset('x_grid', data=np.asarray(mesh.x_grid, dtype=float)) + grp.create_dataset('y_grid', data=np.asarray(mesh.y_grid, dtype=float)) + grp.create_dataset('z_grid', data=np.asarray(mesh.z_grid, dtype=float)) + elif isinstance(mesh, CylindricalMesh): + grp.create_dataset('type', data=np.bytes_('cylindrical')) + grp.create_dataset('r_grid', data=np.asarray(mesh.r_grid, dtype=float)) + grp.create_dataset('phi_grid', data=np.asarray(mesh.phi_grid, dtype=float)) + grp.create_dataset('z_grid', data=np.asarray(mesh.z_grid, dtype=float)) + grp.create_dataset('origin', data=np.asarray(mesh.origin, dtype=float)) + elif isinstance(mesh, SphericalMesh): + grp.create_dataset('type', data=np.bytes_('spherical')) + grp.create_dataset('r_grid', data=np.asarray(mesh.r_grid, dtype=float)) + grp.create_dataset('theta_grid', data=np.asarray(mesh.theta_grid, dtype=float)) + grp.create_dataset('phi_grid', data=np.asarray(mesh.phi_grid, dtype=float)) + grp.create_dataset('origin', data=np.asarray(mesh.origin, dtype=float)) + elif isinstance(mesh, UnstructuredMesh): + raise NotImplementedError( + "Exporting weight windows on an UnstructuredMesh is not yet " + "implemented in Python; use openmc.lib.export_weight_windows()." + ) + else: + raise TypeError(f"Unknown mesh type: {type(mesh).__name__}") + + class WeightWindowsList(list): """A list of WeightWindows objects. @@ -1061,29 +1109,66 @@ def from_wwinp(cls, path: PathLike) -> Self: def export_to_hdf5(self, path: PathLike = 'weight_windows.h5', **init_kwargs): """Write weight windows to an HDF5 file. + Writes the file directly via :mod:`h5py`, mirroring the layout + produced by :func:`openmc.lib.export_weight_windows`. The previous + XML round-trip raised :class:`MemoryError` on multi-GB bound arrays + because of the intermediate ASCII allocation inside lxml. + Parameters ---------- path : PathLike Path to the file to write weight windows to **init_kwargs - Keyword arguments passed to :func:`openmc.lib.init` - + Unused. Retained for backward compatibility (previously forwarded + 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() - # Load the model with openmc.lib and then export it to an HDF5 file - with openmc.lib.TemporarySession(model, **init_kwargs): - openmc.lib.export_weight_windows(path) + with h5py.File(path, 'w') as f: + f.attrs['filetype'] = np.bytes_('weight_windows') + f.attrs['version'] = np.asarray([1, 0], dtype=np.int32) + + meshes_grp = f.create_group('meshes') + wws_grp = f.create_group('weight_windows') + + seen_mesh_ids = set() + ww_ids = [] + for ww in self: + if ww.mesh.id not in seen_mesh_ids: + _write_mesh_group(meshes_grp, ww.mesh) + seen_mesh_ids.add(ww.mesh.id) + + g = wws_grp.create_group(f'weight_windows_{ww.id}') + ww_ids.append(int(ww.id)) + + g.create_dataset('mesh', data=np.int32(ww.mesh.id)) + g.create_dataset('particle_type', + data=np.bytes_(str(ww.particle_type))) + g.create_dataset('energy_bounds', + data=np.asarray(ww.energy_bounds, dtype=float)) + + # 2D (ne, n_voxels) C-contiguous: the C++ reader expects this layout. + ne = ww.lower_ww_bounds.shape[-1] + lo = np.ascontiguousarray(ww.lower_ww_bounds.T).reshape(ne, -1) + up = np.ascontiguousarray(ww.upper_ww_bounds.T).reshape(ne, -1) + g.create_dataset('lower_ww_bounds', data=lo) + g.create_dataset('upper_ww_bounds', data=up) + + g.create_dataset('survival_ratio', + data=float(ww.survival_ratio)) + # max_lower_bound_ratio is read unconditionally by C++; default 1.0 when unset. + mlbr = ww.max_lower_bound_ratio + g.create_dataset( + 'max_lower_bound_ratio', + data=float(mlbr if mlbr is not None else 1.0), + ) + g.create_dataset('max_split', data=np.int32(ww.max_split)) + g.create_dataset('weight_cutoff', + data=float(ww.weight_cutoff)) + + wws_grp.attrs['ids'] = np.asarray(ww_ids, dtype=np.int32) + wws_grp.attrs['n_weight_windows'] = np.int32(len(ww_ids)) + meshes_grp.attrs['ids'] = np.asarray(sorted(seen_mesh_ids), + dtype=np.int32) + meshes_grp.attrs['n_meshes'] = np.int32(len(seen_mesh_ids)) diff --git a/tests/unit_tests/weightwindows/test_ww_list.py b/tests/unit_tests/weightwindows/test_ww_list.py index d148f382a53..17200a656d4 100644 --- a/tests/unit_tests/weightwindows/test_ww_list.py +++ b/tests/unit_tests/weightwindows/test_ww_list.py @@ -1,3 +1,4 @@ +import h5py import openmc @@ -21,3 +22,25 @@ def test_ww_roundtrip(request, run_in_tmpdir): assert ww.max_split == ww_new.max_split assert ww.weight_cutoff == ww_new.weight_cutoff assert ww.mesh.id == ww_new.mesh.id + + +def test_export_hdf5_format(request, run_in_tmpdir): + # C++ openmc_weight_windows_import expects this on-disk layout. + wws = openmc.WeightWindowsList.from_wwinp(request.path.with_name('wwinp_n')) + wws.export_to_hdf5('ww.h5') + with h5py.File('ww.h5') as f: + assert f.attrs['filetype'] == b'weight_windows' + assert list(f.attrs['version']) == [1, 0] + wws_grp = f['weight_windows'] + assert int(wws_grp.attrs['n_weight_windows']) == len(wws) + for ww in wws: + g = wws_grp[f'weight_windows_{ww.id}'] + # 2D shape is required by the C++ tensor::Tensor reader. + assert g['lower_ww_bounds'].ndim == 2 + assert g['lower_ww_bounds'].shape[0] == ww.num_energy_bins + # max_lower_bound_ratio is read unconditionally by C++. + assert 'max_lower_bound_ratio' in g + m_grp = f['meshes'] + for name in m_grp: + assert 'id' in m_grp[name].attrs + assert 'type' in m_grp[name]