Geometry debug function

openmc/model/model.py

@@ -16,6 +16,7 @@
 import lxml.etree as ET
 import numpy as np
 from scipy.optimize import curve_fit
+from scipy import ndimage
 
 import openmc
 import openmc._xml as xml
@@ -2894,6 +2895,278 @@ def _replace_infinity(value):
         # Take a wild guess as to how many rays are needed
         self.settings.particles = 2 * int(max_length)
 
+    @staticmethod
+    def _classify_undefined_regions(cell_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Classify undefined pixels in a 2D cell-ID slice.
+
+        Undefined pixels are identified by the `_NOT_FOUND` sentinel and split
+        into two groups: boundary-connected undefined pixels (`outside`) and
+        non-boundary-connected undefined pixels (`internal`). This classification
+        is based only on connectivity within the sampled pixel grid, so it does
+        not guarantee true geometric exterior/interior classification. To be 
+        reused in the plotter for undefined region visualization.
+
+        Parameters
+        ----------
+        cell_ids : numpy.ndarray
+            Two-dimensional array of cell IDs for a slice, intended to be 
+            gotten from the slice_data function.
+        """
+
+        _NOT_FOUND = -2
+        if cell_ids is None:
+            return None, None, None
+
+        undefined = (cell_ids == _NOT_FOUND)
+
+        # Internal undefined pixels are holes in the defined-pixel mask.
+        internal = ndimage.binary_fill_holes(~undefined) & undefined
+
+        # Anything undefined that is not internal is boundary-connected.
+        outside = undefined & ~internal
+
+        # Guard: undefined regions exist, but none connect to the slice boundary.
+        if undefined.any() and not outside.any():
+            warnings.warn(
+                "Undefined pixels were found, but none are connected to the "
+                "slice boundary. All undefined pixels are being classified as "
+                "internal for this slice. Consider increasing slice resolution."
+            )
+
+        return undefined, outside, internal
+
+    def geometry_debug(
+        self,
+        lower_left,
+        upper_right,
+        n_samples,
+        print_summary=False,
+        **init_kwargs,
+    ):
+        """Sample a 3D region to identify overlap and undefined locations.
+
+        The region between `lower_left` and `upper_right` is sampled on a regular
+        3D grid by taking a sequence of 2D slices in z. Overlap and undefined
+        locations are identified from cells marked with the overlap and undefined
+        sentinels, respectively. A 3D bounding box is returned for each unique
+        overlap pair and for each distinct internal undefined region (found via
+        3D connected-component labeling), in a summary dictionary. This function
+        is meant to be called from an input file on a 3D box encapsulating the
+        entire model.
+
+        Parameters
+        ----------
+        lower_left : Sequence[float]
+            Lower-left corner of the sampled 3D region.
+        upper_right : Sequence[float]
+            Upper-right corner of the sampled 3D region.
+        n_samples : int or Sequence[int]
+            Number of sample points in the x, y, and z directions. If a single
+            integer is given, the value is split into all three directions.
+        print_summary : bool, optional
+            Whether to print a summary of overlap and undefined sample results.
+        **init_kwargs
+            Keyword arguments passed to :meth:`Model.init_lib`.
+        """
+        import openmc.lib
+
+        _OVERLAP = -3
+
+        init_kwargs.setdefault('output', False)
+        init_kwargs.setdefault('args', ['-c'])
+
+        # Accepts 3 separate samples (for x y and z) or just one number
+        if isinstance(n_samples, int):
+            base, extra = divmod(n_samples, 3)
+            nx = base + (1 if extra > 0 else 0)
+            ny = base + (1 if extra > 1 else 0)
+            nz = base
+        else:
+            if len(n_samples) != 3:
+                raise ValueError("n_samples must be an int or a length-3 iterable")
+            nx, ny, nz = n_samples
+
+        nx = int(nx)
+        ny = int(ny)
+        nz = int(nz)
+
+        if nx <= 0 or ny <= 0 or nz <= 0:
+            raise ValueError("All n_samples values must be positive")
+
+        if len(lower_left) != 3:
+            raise ValueError("lower_left must be a length-3 iterable")
+        if len(upper_right) != 3:
+            raise ValueError("upper_right must be a length-3 iterable")
+
+        x0, y0, z0 = lower_left
+        x1, y1, z1 = upper_right
+
+        dz = (z1 - z0) / nz
+
+        u_span = (x1 - x0, 0.0, 0.0)
+        v_span = (0.0, y1 - y0, 0.0)
+
+        # Each unique overlap key (universe, cell1, cell2) gets its own bounding
+        # box, accumulated in world coordinates across all z-slices. Internal
+        # undefined pixels are stacked into a 3D volume and labeled afterwards.
+        overlap_boxes = {}
+        internal_volume = np.zeros((nz, ny, nx), dtype=bool)
+
+        with openmc.lib.TemporarySession(self, **init_kwargs):
+            for k in range(nz):
+                z = z0 + (k + 0.5) * dz
+                origin = ((x0 + x1) / 2.0, (y0 + y1) / 2.0, z)
+
+                geom_data, _ = openmc.lib.slice_data(
+                    origin=origin,
+                    u_span=u_span,
+                    v_span=v_span,
+                    pixels=(nx, ny),
+                    show_overlaps=True,
+                    level=-1,
+                    include_properties=False,
+                )
+
+                cell_ids = geom_data[:, :, 0]
+
+                overlap_data = openmc.lib.slice_data_overlap_info()
+
+                # Extend the bounding box for each unique overlap key.
+                for overlap_idx, key in enumerate(overlap_data):
+                    encoded_id = _OVERLAP - overlap_idx - 1
+                    pix = np.argwhere(cell_ids == encoded_id)
+                    if pix.size == 0:
+                        continue
+
+                    key_t = tuple(int(v) for v in key)
+                    xc = x0 + (pix[:, 1] + 0.5) * (x1 - x0) / nx
+                    yc = y1 - (pix[:, 0] + 0.5) * (y1 - y0) / ny
+
+                    box = overlap_boxes.get(key_t)
+                    if box is None:
+                        overlap_boxes[key_t] = {
+                            "key": key_t,
+                            "xmin": float(xc.min()), "xmax": float(xc.max()),
+                            "ymin": float(yc.min()), "ymax": float(yc.max()),
+                            "zmin": float(z), "zmax": float(z),
+                        }
+                    else:
+                        box["xmin"] = min(box["xmin"], float(xc.min()))
+                        box["xmax"] = max(box["xmax"], float(xc.max()))
+                        box["ymin"] = min(box["ymin"], float(yc.min()))
+                        box["ymax"] = max(box["ymax"], float(yc.max()))
+                        box["zmin"] = min(box["zmin"], float(z))
+                        box["zmax"] = max(box["zmax"], float(z))
+
+                _, _, internal = Model._classify_undefined_regions(cell_ids)
+
+                internal_volume[k] = internal
+
+        overlap_boxes = list(overlap_boxes.values())
+
+        # Overlaps are not flagged for resolution: whether an overlap exists and
+        # which cells collide is determined by the (universe, cell1, cell2) key
+
+        # Label spatially-connected undefined regions in 3D and build a
+        # world-coordinate bounding box for each connected component. A feature 
+        # thinner than the sample spacing can rasterize with breaks and split 
+        # into several regions, so give a suggestion to the user to increase n_samples.
+
+        undefined_boxes = []
+        if internal_volume.any():
+            structure = ndimage.generate_binary_structure(3, 1)  # face connectivity
+            labeled, _ = ndimage.label(internal_volume, structure=structure)
+            for region_id, sl in enumerate(ndimage.find_objects(labeled), start=1):
+                if sl is None:
+                    continue
+                kz, ky, kx = sl  # slice objects over the (z, y, x) index axes
+
+                # Voxel-edge world extents (start inclusive, stop exclusive)
+                x_lo = x0 + kx.start * (x1 - x0) / nx
+                x_hi = x0 + kx.stop * (x1 - x0) / nx
+                # The y (row) axis is flipped in world coordinates (row 0 == y1)
+                y_hi = y1 - ky.start * (y1 - y0) / ny
+                y_lo = y1 - ky.stop * (y1 - y0) / ny
+                z_lo = z0 + kz.start * dz
+                z_hi = z0 + kz.stop * dz
+
+                # Local-thickness test: a bounding box is misleading for thin
+                # curved shells (e.g. an annular gap whose bbox is large but
+                # which is only ~1 voxel thick radially). Erode the region's
+                # voxel mask; if erosion empties it, the region is nowhere
+                # thicker than ~2 voxels and is under-resolved.
+                mask = (labeled[sl] == region_id)
+                under_resolved = not ndimage.binary_erosion(mask).any()
+
+                undefined_boxes.append({
+                    "xmin": float(min(x_lo, x_hi)), "xmax": float(max(x_lo, x_hi)),
+                    "ymin": float(min(y_lo, y_hi)), "ymax": float(max(y_lo, y_hi)),
+                    "zmin": float(z_lo), "zmax": float(z_hi),
+                    "under_resolved": bool(under_resolved),
+                })
+
+        # Round coordinates to keep the reported boxes readable.
+        for box in overlap_boxes + undefined_boxes:
+            for coord in ("xmin", "xmax", "ymin", "ymax", "zmin", "zmax"):
+                box[coord] = round(box[coord], 4)
+
+        under_resolved = any(b["under_resolved"] for b in undefined_boxes)
+
+        result = {
+            "overlap_boxes": overlap_boxes,
+            "undefined_boxes": undefined_boxes,
+            "n_overlaps": len(overlap_boxes),
+            "n_undefined_regions": len(undefined_boxes),
+            "under_resolved": under_resolved,
+        }
+
+        if under_resolved:
+            n_un = sum(b["under_resolved"] for b in undefined_boxes)
+            warnings.warn(
+                f"Sampling resolution may be insufficient: {n_un} undefined "
+                "region(s) are resolved by <= 2 voxels across their thinnest "
+                "dimension, so they may be fragmented. "
+                "Consider increasing n_samples."
+            )
+
+        if print_summary:
+            print("Geometry debug summary:")
+
+            if result["overlap_boxes"]:
+                print(f"  Overlaps found: {result['n_overlaps']}")
+                for box in result["overlap_boxes"]:
+                    print(
+                        f"    cells {box['key']}: "
+                        f"x[{box['xmin']:.4g}, {box['xmax']:.4g}] "
+                        f"y[{box['ymin']:.4g}, {box['ymax']:.4g}] "
+                        f"z[{box['zmin']:.4g}, {box['zmax']:.4g}]"
+                    )
+            else:
+                print("  Overlap bounding boxes: None")
+
+            if result["undefined_boxes"]:
+                print(f"  Undefined regions found: {result['n_undefined_regions']}")
+                for i, box in enumerate(result["undefined_boxes"], start=1):
+                    flag = "  [under-resolved]" if box["under_resolved"] else ""
+                    print(
+                        f"    region {i}: "
+                        f"x[{box['xmin']:.4g}, {box['xmax']:.4g}] "
+                        f"y[{box['ymin']:.4g}, {box['ymax']:.4g}] "
+                        f"z[{box['zmin']:.4g}, {box['zmax']:.4g}]{flag}"
+                    )
+            else:
+                print("  Undefined bounding boxes: None")
+
+            if result["under_resolved"]:
+                print(
+                    "  WARNING: some undefined regions are resolved by <= 2 "
+                    "voxels across their thinnest dimension and may be "
+                    "fragmented or missed; increase n_samples so thin features "
+                    "span at least 3 voxels."
+                )
+
+        return result
+
     def keff_search(
         self,
         func: ModelModifier,