Make distribcell setup scale with the number of universes

[GuySten]

Description

Two steps of geometry setup scale with the number of universes times the size of the geometry:

1. Counting universe instances searches the geometry from the root once per universe.
2. Distribcell offset tables store, for every fill cell and every lattice tile, one offset per map, where a map is a universe containing distributed cells. By default every universe containing a material cell is a map.

Models with many universes hit both. One example is a 160 cm TRISO fuel compact where each tile of a TRISO search lattice has its own universe: 20,485 maps and 239,490 fill cells. On develop, the offset tables for this model need 21.3 GB (19.6 GB for cells, 1.7 GB for the lattice), while the rest of the simulation needs under 400 MB. With material_cell_offsets = False, the model fits in memory, but setup takes 19 minutes, almost all of it counting instances.

This PR has one commit for each step:

1. Count universe instances in one pass. Universes are ordered so that each comes after the universes containing it, and counts are propagated from the root in that order.
2. Store distribcell offsets only where they can be read. An offset is only read when the map's universe is inside the cell's fill or the tile's universe. Each universe and lattice now keeps the sorted list of maps it contains, and each fill cell and tile stores offsets for those maps only, through a new DistribcellOffsets class. The offsets of a universe's fill cells, and of a lattice's tiles, are stored together in one array. Maps are numbered in post-order, so the maps in a universe are mostly consecutive and a lookup is usually one range check, with a binary search as fallback. The offsets are found in one visit of each universe and lattice from the root, instead of one pass over the geometry per map.

The stored values are the same as before, so instance numbers don't change. Callers of Cell::offset_ and Lattice::offset() are unchanged. distribcell_path_inner now skips cells and tiles that don't contain the map; they contain no instances of it, so the path it finds is the same.

Results

160 cm TRISO compact, default settings, single thread:

	Peak memory	Setup time
develop	runs out of memory (offset tables alone: 21.3 GB)	—
develop, material_cell_offsets = False	387 MB	1,130 s
This PR, default settings	392 MB	105 s

k-effective is identical. On a 10 cm slice of the same compact, setup goes from 11–13.5 s to 6.8 s.

Transport cost

Share of profiled run time in the instance lookup (cell_instance_at_level and the offset accessors), single thread, two or three runs each:

Model	develop	This PR
PWR core (pins in a core lattice)	1.47–1.56%	1.67–2.02%
TRISO compact, 10 cm slice	3.94–4.17%	2.91–2.96%
ATR (flat CSG, 5,968 cells)	0.06–0.08%	0.06–0.07%

On PWR, wall-clock transport time over 6 alternating runs is within noise of develop (medians differ by 0.9%, fastest runs by −2.4%).

Testing

• Instance numbers compared with develop on 17 geometries: rect and hex lattices, outer universes, nested lattices, a synthetic model with shared universes at several depths, TRISO, PWR and ATR. slice_data cell instances at every geometry level and DistribcellFilter bins, on slices along three axes: 265,791 distinct (cell, instance) pairs, all identical.
• tallies.out from short runs with DistribcellFilter tallies on 13 models: identical to develop, including about 10,000 distributed cell labels.
• Added tests/unit_tests/test_universe_instances.py for the instance counts.
• Full test suite passes locally. The only errors are tests that need ENDF data or dlopen, which fail the same way on develop.

Checklist

• I have performed a self-review of my own code
• I have run clang-format (version 18) on any C++ source files (if applicable)
• I have followed the style guidelines for Python source files (if applicable)
• I have made corresponding changes to the documentation (if applicable)
• I have added tests that prove my fix is effective or that my feature works (if applicable)

[GuySten]

@nuclearkevin this PR came out of trying the full 160 cm HTGR compact from your delta tracking work (the one with the TRISO search lattice) on develop. With default settings it can't run: the distribcell offset tables alone need about 21 GB, because each search-lattice tile has its own universe (about 20,500 of them) and the tables have one entry per universe for every fill cell. With material_cell_offsets = False it fits, but setup takes about 19 minutes, almost all of it in count_universe_instances.

With this PR, on my reconstruction of your compact (single thread, default settings):

	Peak memory	Setup time
develop, default settings	doesn't fit (21.3 GB of offset tables)	—
develop, material_cell_offsets = False	387 MB	1,130 s
This PR, default settings	392 MB	105 s

Instance numbers are unchanged, and transport on the compact is the same or slightly faster.

[nuclearkevin]

Thanks for looking into this @GuySten! I wonder if this will also decrease the time spent constructing the majorant in the delta tracking branch for TRISO problems, which has to do a forall cell instances lookup to find the maximum cell density per material. I'll try this out and see how it performs there.

Edit: this also reminds me, I need to look at the 2-component combined estimator PR you put in. I'll get to that this weekend, sorry about the wait!