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

.github/workflows/ci.yml

@@ -177,9 +177,9 @@ jobs:
 
       - name: Setup tmate debug session
         continue-on-error: true
-        if: ${{ contains(env.COMMIT_MESSAGE, '[gha-debug]') }}
+        if: ${{ failure() && contains(env.COMMIT_MESSAGE, '[gha-debug]') }}
         uses: mxschmitt/action-tmate@v3
-        timeout-minutes: 10
+        timeout-minutes: 100
 
       - name: Generate C++ coverage (gcovr)
         shell: bash

include/openmc/capi.h

@@ -1,6 +1,8 @@
 #ifndef OPENMC_CAPI_H
 #define OPENMC_CAPI_H
 
+#include "hdf5.h"
+
 #include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
@@ -74,6 +76,7 @@ int openmc_get_n_batches(int* n_batches, bool get_max_batches);
 int openmc_get_nuclide_index(const char name[], int* index);
 int openmc_add_unstructured_mesh(
   const char filename[], const char library[], int* id);
+int openmc_unstructured_mesh_export_hdf5(int32_t index, hid_t mesh_group);
 int64_t openmc_get_seed();
 uint64_t openmc_get_stride();
 int openmc_get_tally_index(int32_t id, int32_t* index);

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,10 +1,11 @@
 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
 from weakref import WeakValueDictionary
 
+import h5py
 import numpy as np
 from numpy.ctypeslib import as_array
 
@@ -18,10 +19,15 @@
 
 __all__ = [
     'Mesh', 'RegularMesh', 'RectilinearMesh', 'CylindricalMesh',
-    'SphericalMesh', 'UnstructuredMesh', 'meshes', 'MeshMaterialVolumes'
+    'SphericalMesh', 'UnstructuredMesh', 'meshes', 'MeshMaterialVolumes', 'export_unstructured_mesh'
 ]
 
 
+if (h5py.h5.get_libversion() >= (1, 10, 0)):
+    c_hid_t = c_int64
+else:
+    c_hid_t = c_int32
+
 arr_2d_int32 = np.ctypeslib.ndpointer(dtype=np.int32, ndim=2, flags='CONTIGUOUS')
 arr_2d_double = np.ctypeslib.ndpointer(dtype=np.double, ndim=2, flags='CONTIGUOUS')
 
@@ -108,6 +114,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_hid_t]
+_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 +750,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

@@ -304,6 +304,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
 
@@ -1229,6 +1253,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,
@@ -1706,6 +1743,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
@@ -2102,6 +2150,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,
@@ -2474,6 +2535,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,
@@ -3290,6 +3364,22 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    def to_hdf5(self, group: h5py.Group):
+        import openmc.lib
+
+        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, 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,21 @@ 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 (!mpi::master)
+    return 0;
+
+  if (int err = check_mesh_type<UnstructuredMesh>(index))
+    return err;
+  UnstructuredMesh* m =
+    dynamic_cast<UnstructuredMesh*>(model::meshes[index].get());
+
+  m->to_hdf5(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()