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

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,20 @@ 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: trusted dtype skips the per-element scan (see Notes).
+    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/mesh.py

@@ -306,6 +306,29 @@ def from_hdf5(cls, group: h5py.Group):
         else:
             raise ValueError('Unrecognized mesh type: "' + mesh_type + '"')
 
+    def to_hdf5(self, group: h5py.Group) -> h5py.Group:
+        """Write this mesh into *group* as a subgroup named ``mesh <id>``.
+
+        Subclasses override this method to call ``super().to_hdf5(group)``,
+        write a ``type`` dataset, and append type-specific grid data.
+
+        .. versionadded:: 0.15.4
+
+        Parameters
+        ----------
+        group : h5py.Group
+            Parent HDF5 group (typically ``/meshes``).
+
+        Returns
+        -------
+        mesh_group : h5py.Group
+            The created ``mesh <id>`` subgroup, ready for subclass extension.
+        """
+        mesh_group = group.create_group(f'mesh {self.id}')
+        mesh_group.attrs['id'] = np.int32(self.id)
+        mesh_group.create_dataset('name', data=np.bytes_(self.name or ''))
+        return mesh_group
+
     def to_xml_element(self):
         """Return XML representation of the mesh
 
@@ -1231,6 +1254,18 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        mesh_group = super().to_hdf5(group)
+        mesh_group.create_dataset('type', data=np.bytes_('regular'))
+        mesh_group.create_dataset(
+            'dimension', data=np.asarray(self.dimension, dtype=np.int32))
+        mesh_group.create_dataset(
+            'lower_left', data=np.asarray(self.lower_left, dtype=float))
+        mesh_group.create_dataset(
+            'upper_right', data=np.asarray(self.upper_right, dtype=float))
+        mesh_group.create_dataset(
+            'width', data=np.asarray(self.width, dtype=float))
+
     @classmethod
     def from_rect_lattice(
         cls,
@@ -1708,6 +1743,16 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        mesh_group = super().to_hdf5(group)
+        mesh_group.create_dataset('type', data=np.bytes_('rectilinear'))
+        mesh_group.create_dataset(
+            'x_grid', data=np.asarray(self.x_grid, dtype=float))
+        mesh_group.create_dataset(
+            'y_grid', data=np.asarray(self.y_grid, dtype=float))
+        mesh_group.create_dataset(
+            'z_grid', data=np.asarray(self.z_grid, dtype=float))
+
     @classmethod
     def from_xml_element(cls, elem: ET.Element):
         """Generate a rectilinear mesh from an XML element
@@ -2104,6 +2149,18 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        mesh_group = super().to_hdf5(group)
+        mesh_group.create_dataset('type', data=np.bytes_('cylindrical'))
+        mesh_group.create_dataset(
+            'r_grid', data=np.asarray(self.r_grid, dtype=float))
+        mesh_group.create_dataset(
+            'phi_grid', data=np.asarray(self.phi_grid, dtype=float))
+        mesh_group.create_dataset(
+            'z_grid', data=np.asarray(self.z_grid, dtype=float))
+        mesh_group.create_dataset(
+            'origin', data=np.asarray(self.origin, dtype=float))
+
     @classmethod
     def from_bounding_box(
         cls,
@@ -2476,6 +2533,18 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        mesh_group = super().to_hdf5(group)
+        mesh_group.create_dataset('type', data=np.bytes_('spherical'))
+        mesh_group.create_dataset(
+            'r_grid', data=np.asarray(self.r_grid, dtype=float))
+        mesh_group.create_dataset(
+            'theta_grid', data=np.asarray(self.theta_grid, dtype=float))
+        mesh_group.create_dataset(
+            'phi_grid', data=np.asarray(self.phi_grid, dtype=float))
+        mesh_group.create_dataset(
+            'origin', data=np.asarray(self.origin, dtype=float))
+
     @classmethod
     def from_bounding_box(
         cls,
@@ -3220,6 +3289,14 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        # Raise before super() so no half-built 'mesh <id>' group is left on disk.
+        raise NotImplementedError(
+            "UnstructuredMesh.to_hdf5 is not implemented in Python. "
+            "Use openmc.lib.export_weight_windows() to export weight "
+            "windows on unstructured meshes."
+        )
+
     def to_xml_element(self):
         """Return XML representation of the mesh
 

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
@@ -1063,29 +1061,83 @@ 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)
+        # Any unstructured mesh forces the whole list onto the lib fallback.
+        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:
+                    ww.mesh.to_hdf5(meshes_grp)
+                    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]