perf: use compact array format for TRISO particle HDF5 I/O

Replace per-particle HDF5 group writing with a compact array format
that stores only (x, y, z, radius) + IDs per TRISO particle. This
eliminates the massive HDF5 group metadata overhead and supports
multiple TRISO groups with different layer structures or background
materials.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: cn-skywalker

openmc/summary.py

@@ -112,6 +112,10 @@ def _read_geometry(self):
             if "dagmc" not in self._f['geometry'].attrs.keys():
                 self._read_surfaces()
                 cell_fills = self._read_cells()
+                # Read compact TRISO particle data if present
+                if 'geometry/triso_particles' in self._f:
+                    triso_fills = self._read_triso_particles()
+                    cell_fills.update(triso_fills)
                 self._read_universes()
                 self._read_lattices()
                 self._finalize_geometry(cell_fills)
@@ -194,6 +198,69 @@ def _read_cells(self):
 
         return cell_fills
 
+    def _read_triso_particles(self):
+        """Read compact TRISO particle data from HDF5 and reconstruct objects.
+
+        Returns
+        -------
+        dict
+            Cell fill information for all TRISO and background cells
+        """
+        triso_grp = self._f['geometry/triso_particles']
+        n_triso = int(triso_grp.attrs['n_triso'])
+        if n_triso == 0:
+            return {}
+
+        # Batch-read all arrays
+        positions = triso_grp['positions'][()]       # (N, 4) array
+        surface_ids = triso_grp['surface_ids'][()]    # (N,) array
+        cell_ids = triso_grp['cell_ids'][()]          # (N,) array
+        fill_universes = triso_grp['fill_universes'][()]  # (N,) array
+        group_offsets = triso_grp['group_offsets'][()]    # (n_groups+1,) array
+
+        cell_fills = {}
+
+        # Create all TRISO Sphere surfaces in batch
+        for i in range(n_triso):
+            x0, y0, z0, r = positions[i]
+            sid = int(surface_ids[i])
+            surf = openmc.Sphere(x0=float(x0), y0=float(y0),
+                                 z0=float(z0), r=float(r),
+                                 surface_id=sid)
+            self._fast_surfaces[sid] = surf
+
+        # Create all TRISO Cell objects in batch
+        for i in range(n_triso):
+            cid = int(cell_ids[i])
+            sid = int(surface_ids[i])
+            fill_univ_id = int(fill_universes[i])
+            x0, y0, z0, _ = positions[i]
+
+            cell = openmc.Cell(cell_id=cid)
+            cell.region = -self._fast_surfaces[sid]
+            cell.translation = np.array([x0, y0, z0], dtype=np.float64)
+            cell_fills[cid] = ('universe', fill_univ_id)
+            self._fast_cells[cid] = cell
+
+        # Create background cells for each group
+        n_groups = len(group_offsets) - 1
+        groups_grp = triso_grp['groups']
+        for g in range(n_groups):
+            grp = groups_grp[str(g)]
+            bg_cell_id = int(grp['background_cell_id'][()])
+            bg_material_id = int(grp['background_material_id'][()])
+            bg_universe_id = int(grp['background_universe_id'][()])
+
+            bg_cell = openmc.Cell(cell_id=bg_cell_id)
+            bg_cell.virtual_lattice = True
+            bg_cell.lower_left = list(grp['vl_lower_left'][()])
+            bg_cell.pitch = list(grp['vl_pitch'][()])
+            bg_cell.shape = list(grp['vl_shape'][()])
+            cell_fills[bg_cell_id] = ('material', bg_material_id)
+            self._fast_cells[bg_cell_id] = bg_cell
+
+        return cell_fills
+
     def _read_universes(self):
         for group in self._f['geometry/universes'].values():
             geom_type = group.get('geom_type')

src/summary.cpp

@@ -1,6 +1,7 @@
 #include "openmc/summary.h"
 
 #include <fmt/core.h>
+#include <map>
 
 #include "openmc/capi.h"
 #include "openmc/cell.h"
@@ -93,25 +94,139 @@ void write_nuclides(hid_t file)
   close_group(macro_group);
 }
 
+void write_triso_particles(hid_t geom_group)
+{
+  // Collect all TRISO surfaces and group by background cell ID
+  std::map<int32_t, vector<int32_t>> bg_to_surfs;
+  for (int i = 0; i < model::surfaces.size(); ++i) {
+    const auto& surf = model::surfaces[i];
+    if (surf->is_triso_surface_) {
+      bg_to_surfs[surf->triso_base_index_].push_back(i);
+    }
+  }
+
+  int n_groups = bg_to_surfs.size();
+  if (n_groups == 0)
+    return;
+
+  // Count total TRISO particles
+  int n_triso = 0;
+  for (const auto& [bg_id, surfs] : bg_to_surfs) {
+    n_triso += surfs.size();
+  }
+
+  // Build flat arrays
+  vector<double> positions(n_triso * 4);
+  vector<int32_t> surface_ids(n_triso);
+  vector<int32_t> cell_ids(n_triso);
+  vector<int32_t> fill_universes(n_triso);
+  vector<int32_t> group_offsets(n_groups + 1);
+
+  int idx = 0;
+  int g = 0;
+  group_offsets[0] = 0;
+
+  for (const auto& [bg_cell_id, surfs] : bg_to_surfs) {
+    for (const auto& i_surf : surfs) {
+      const auto& surf = model::surfaces[i_surf];
+      vector<double> center = surf->get_center();
+      double radius = surf->get_radius();
+      positions[idx * 4 + 0] = center[0];
+      positions[idx * 4 + 1] = center[1];
+      positions[idx * 4 + 2] = center[2];
+      positions[idx * 4 + 3] = radius;
+      surface_ids[idx] = surf->id_;
+      // Get TRISO particle cell via triso_particle_index_
+      int32_t i_cell = model::cell_map.at(surf->triso_particle_index_);
+      const auto& cell = model::cells[i_cell];
+      cell_ids[idx] = cell->id_;
+      fill_universes[idx] = model::universes[cell->fill_]->id_;
+      ++idx;
+    }
+    group_offsets[g + 1] = idx;
+    ++g;
+  }
+
+  // Write compact TRISO data to HDF5
+  hid_t triso_group = create_group(geom_group, "triso_particles");
+  write_attribute(triso_group, "n_groups", n_groups);
+  write_attribute(triso_group, "n_triso", n_triso);
+
+  // Write positions as 2D dataset [n_triso, 4]
+  hsize_t pos_dims[] = {static_cast<hsize_t>(n_triso), 4};
+  write_dataset_lowlevel(triso_group, 2, pos_dims, "positions",
+    H5TypeMap<double>::type_id, H5S_ALL, false, positions.data());
+
+  write_dataset(triso_group, "surface_ids", surface_ids);
+  write_dataset(triso_group, "cell_ids", cell_ids);
+  write_dataset(triso_group, "fill_universes", fill_universes);
+  write_dataset(triso_group, "group_offsets", group_offsets);
+
+  // Write per-group metadata (background cell info)
+  hid_t groups_group = create_group(triso_group, "groups");
+  g = 0;
+  for (const auto& [bg_cell_id, surfs] : bg_to_surfs) {
+    auto grp = create_group(groups_group, std::to_string(g));
+
+    int32_t i_bg_cell = model::cell_map.at(bg_cell_id);
+    const auto& bg_cell = model::cells[i_bg_cell];
+
+    write_dataset(grp, "background_cell_id", bg_cell->id_);
+    write_dataset(grp, "background_material_id",
+      model::materials[bg_cell->material_[0]]->id_);
+    write_dataset(grp, "background_universe_id",
+      model::universes[bg_cell->universe_]->id_);
+    write_dataset(grp, "vl_lower_left", bg_cell->vl_lower_left_);
+    write_dataset(grp, "vl_pitch", bg_cell->vl_pitch_);
+    write_dataset(grp, "vl_shape", bg_cell->vl_shape_);
+
+    close_group(grp);
+    ++g;
+  }
+  close_group(groups_group);
+  close_group(triso_group);
+}
+
 void write_geometry(hid_t file)
 {
   auto geom_group = create_group(file, "geometry");
 
-  write_attribute(geom_group, "n_cells", model::cells.size());
-  write_attribute(geom_group, "n_surfaces", model::surfaces.size());
+  // Count non-TRISO cells and surfaces for accurate counts
+  int n_cells = 0;
+  int n_surfaces = 0;
+  for (const auto& c : model::cells) {
+    if (!c->triso_particle_ && !c->virtual_lattice_)
+      ++n_cells;
+  }
+  for (const auto& surf : model::surfaces) {
+    if (!surf->is_triso_surface_)
+      ++n_surfaces;
+  }
+
+  write_attribute(geom_group, "n_cells", n_cells);
+  write_attribute(geom_group, "n_surfaces", n_surfaces);
   write_attribute(geom_group, "n_universes", model::universes.size());
   write_attribute(geom_group, "n_lattices", model::lattices.size());
 
   auto cells_group = create_group(geom_group, "cells");
-  for (const auto& c : model::cells)
+  for (const auto& c : model::cells) {
+    if (c->triso_particle_ || c->virtual_lattice_)
+      continue;
     c->to_hdf5(cells_group);
+  }
   close_group(cells_group);
 
   auto surfaces_group = create_group(geom_group, "surfaces");
-  for (const auto& surf : model::surfaces)
+  for (const auto& surf : model::surfaces) {
+    if (surf->is_triso_surface_)
+      continue;
     surf->to_hdf5(surfaces_group);
+  }
   close_group(surfaces_group);
 
+  // Write compact TRISO particle data
+  write_triso_particles(geom_group);
+
   auto universes_group = create_group(geom_group, "universes");
   for (const auto& u : model::universes)
     u->to_hdf5(universes_group);