wwinp files: Fix MemoryError in WeightWindowsList.export_to_hdf5 and speed up from_wwinp

openmc/checkvalue.py

@@ -1,6 +1,7 @@
 import copy
 import os
 from collections.abc import Iterable
+from numbers import Real
 
 import numpy as np
 
@@ -80,7 +81,21 @@ 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
+    -----
+    For numpy float/complex ndarrays whose ``ndim`` is within
+    ``[min_depth, max_depth]`` and whose ``expected_type`` is ``Real``,
+    ``float``, or ``complex``, the dtype guarantees the element type and the
+    per-element scan is skipped for faster processing.
     """
+    # Fast path: float/complex ndarrays whose dtype already guarantees the
+    # expected_type can skip the per-element scan.
+    if (isinstance(value, np.ndarray) and value.dtype.kind in 'fc'
+            and min_depth <= value.ndim <= max_depth
+            and expected_type in (Real, float, complex)):
+        return
+
     # Initialize the tree at the very first item.
     tree = [value]
     index = [0]

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
@@ -140,9 +143,9 @@ 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 = (
+                np.asarray(self.lower_ww_bounds) * upper_bound_ratio
+            )
 
         if upper_ww_bounds is not None:
             self.upper_ww_bounds = upper_ww_bounds
@@ -821,6 +824,46 @@ 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 <id>`` 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))
+    else:
+        raise TypeError(f"Unknown mesh type: {type(mesh).__name__}")
+
+
 class WeightWindowsList(list):
     """A list of WeightWindows objects.
 
@@ -1063,29 +1106,82 @@ 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.
 
+        Structured-mesh weight windows are written directly via :mod:`h5py`,
+        avoiding the multi-GB ASCII intermediate that previously caused
+        :class:`MemoryError` on large wwinp inputs.
+
+        Weight windows on an :class:`~openmc.UnstructuredMesh` require
+        LibMesh or MOAB (loaded by :func:`openmc.lib.init`) to materialize
+        the mesh's vertex and connectivity data; for those, this method
+        falls back to :func:`openmc.lib.export_weight_windows`.
+
         Parameters
         ----------
         path : PathLike
             Path to the file to write weight windows to
         **init_kwargs
-            Keyword arguments passed to :func:`openmc.lib.init`
-
+            Forwarded to :func:`openmc.lib.init` when the UnstructuredMesh
+            fallback path is used. Unused for the direct h5py path.
         """
-        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)
+        if any(isinstance(ww.mesh, UnstructuredMesh) for ww in self):
+            import openmc.lib
+            model = openmc.Model()
+            sph = openmc.Sphere(boundary_type='vacuum')
+            model.geometry = openmc.Geometry([openmc.Cell(region=-sph)])
+            model.settings.weight_windows = self
+            model.settings.particles = 100
+            model.settings.batches = 1
+            with openmc.lib.TemporarySession(model, **init_kwargs):
+                openmc.lib.export_weight_windows(path)
+            return
+
+        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))

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<double> 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]