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 ndarrays whose ``ndim`` is within
+    ``[min_depth, max_depth]`` and whose ``expected_type`` is ``Real`` or
+    ``float``, the dtype guarantees the element type and the per-element scan
+    is skipped for faster processing.
     """
+    # Fast path: trusted float dtype skips the per-element scan (see Notes).
+    if (isinstance(value, np.ndarray) and value.dtype.kind == 'f'
+            and min_depth <= value.ndim <= max_depth
+            and expected_type in (Real, float)):
+        return
+
     # Initialize the tree at the very first item.
     tree = [value]
     index = [0]

openmc/lib/mesh.py

@@ -1,5 +1,5 @@
 from collections.abc import Mapping, Sequence
-from ctypes import (c_int, c_int32, c_char_p, c_double, POINTER, c_void_p,
+from ctypes import (c_int, c_int32, c_int64, c_char_p, c_double, POINTER, c_void_p,
                     create_string_buffer, c_size_t)
 from math import sqrt
 import sys
@@ -18,7 +18,7 @@
 
 __all__ = [
     'Mesh', 'RegularMesh', 'RectilinearMesh', 'CylindricalMesh',
-    'SphericalMesh', 'UnstructuredMesh', 'meshes', 'MeshMaterialVolumes'
+    'SphericalMesh', 'UnstructuredMesh', 'meshes', 'MeshMaterialVolumes', 'export_unstructured_mesh'
 ]
 
 
@@ -108,6 +108,10 @@
 _dll.openmc_spherical_mesh_set_grid.restype = c_int
 _dll.openmc_spherical_mesh_set_grid.errcheck = _error_handler
 
+_dll.openmc_unstructured_mesh_export_hdf5.argtypes = [c_int32, c_int64]
+_dll.openmc_unstructured_mesh_export_hdf5.restype = c_int
+_dll.openmc_unstructured_mesh_export_hdf5.errcheck = _error_handler
+
 
 class Mesh(_FortranObjectWithID):
     """Base class to represent mesh objects
@@ -740,6 +744,12 @@ class UnstructuredMesh(Mesh):
 }
 
 
+def export_unstructured_mesh(mesh, group):
+    index = c_int32()
+    _dll.openmc_get_mesh_index(mesh.id, index)
+    _dll.openmc_unstructured_mesh_export_hdf5(index.value, int(group.id.id))
+
+
 def _get_mesh(index):
     mesh_type = create_string_buffer(20)
     _dll.openmc_mesh_get_type(index, mesh_type)

openmc/mesh.py

@@ -306,6 +306,30 @@ def from_hdf5(cls, group: h5py.Group):
         else:
             raise ValueError('Unrecognized mesh type: "' + mesh_type + '"')
 
+    @abstractmethod
+    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 +1255,19 @@ 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))
+        return mesh_group
+
     @classmethod
     def from_rect_lattice(
         cls,
@@ -1708,6 +1745,17 @@ 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))
+        return mesh_group
+
     @classmethod
     def from_xml_element(cls, elem: ET.Element):
         """Generate a rectilinear mesh from an XML element
@@ -2104,6 +2152,19 @@ 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))
+        return mesh_group
+
     @classmethod
     def from_bounding_box(
         cls,
@@ -2476,6 +2537,19 @@ 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))
+        return mesh_group
+
     @classmethod
     def from_bounding_box(
         cls,
@@ -3220,6 +3294,24 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        import openmc.lib
+
+        mesh_group = super().to_hdf5(group)
+        model = openmc.Model()
+        sph = openmc.Sphere(boundary_type='vacuum')
+        model.geometry = openmc.Geometry([openmc.Cell(region=-sph)])
+        model.settings.particles = 100
+        model.settings.batches = 1
+        wwg = openmc.WeightWindowGenerator(
+              method='magic',
+              mesh=self,
+        )
+        model.settings.weight_window_generators = [wwg] 
+        with openmc.lib.TemporarySession(model):
+            openmc.lib.export_unstructured_mesh(self, mesh_group)
+        return mesh_group
+
     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
@@ -1060,32 +1058,70 @@ def from_wwinp(cls, path: PathLike) -> Self:
 
         return wws
 
-    def export_to_hdf5(self, path: PathLike = 'weight_windows.h5', **init_kwargs):
+    def export_to_hdf5(self, path: PathLike = 'weight_windows.h5'):
         """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
+        uses the c-api under the hood.
+
         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()
 
-        # 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:
+                    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))

src/mesh.cpp

@@ -2843,6 +2843,27 @@ extern "C" int openmc_spherical_mesh_set_grid(int32_t index,
     index, grid_x, nx, grid_y, ny, grid_z, nz);
 }
 
+extern "C" int openmc_unstructured_mesh_export_hdf5(
+  int32_t index, hid_t mesh_group)
+{
+  if (int err = check_mesh_type<UnstructuredMesh>(index))
+    return err;
+  UnstructuredMesh* m =
+    dynamic_cast<UnstructuredMesh*>(model::meshes[index].get());
+
+  // Write mesh type
+  write_dataset(mesh_group, "type", m->get_mesh_type());
+
+  // Write mesh ID
+  write_attribute(mesh_group, "id", m->id_);
+
+  // Write mesh name
+  write_dataset(mesh_group, "name", m->name_);
+
+  m->to_hdf5_inner(mesh_group);
+  return 0;
+}
+
 #ifdef OPENMC_DAGMC_ENABLED
 
 const std::string MOABMesh::mesh_lib_type = "moab";

src/weight_windows.cpp

@@ -1036,6 +1036,7 @@ void free_memory_weight_windows()
 {
   variance_reduction::ww_map.clear();
   variance_reduction::weight_windows.clear();
+  variance_reduction::weight_windows_generators.clear();
 }
 
 void finalize_variance_reduction()

tests/unit_tests/weightwindows/test_ww_list.py

@@ -1,3 +1,6 @@
+import h5py
+import numpy as np
+import pytest
 import openmc
 
 
@@ -21,3 +24,48 @@ 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]
+
+
+@pytest.mark.parametrize('library', ('libmesh', 'moab'))
+def test_export_hdf5_unstructured_mesh(request, run_in_tmpdir, library):
+    # UnstructuredMesh can't be serialized from pure Python; export_to_hdf5
+    # routes it through openmc.lib.export_weight_windows (a live session).
+    if library == 'libmesh' and not openmc.lib._libmesh_enabled():
+        pytest.skip('LibMesh not enabled in this build.')
+    if library == 'moab' and not openmc.lib._dagmc_enabled():
+        pytest.skip('DAGMC (and MOAB) not enabled in this build.')
+
+    mesh = openmc.UnstructuredMesh(
+        str(request.path.with_name('test_mesh_tets.exo')), library)
+    ww = openmc.WeightWindows(mesh, np.ones((12_000,)), upper_bound_ratio=5.0)
+    openmc.WeightWindowsList([ww]).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]
+        assert int(f['weight_windows'].attrs['n_weight_windows']) == 1
+        m_grp = f['meshes'][f'mesh {mesh.id}']
+        assert m_grp['type'][()] == b'unstructured'
+        assert 'filename' in m_grp