Allow MOAB k-d tree to be configured

[paulromano]

Description

If you have unstructured mesh tallies based on MOAB, OpenMC will utilize an adaptive k-d tree (also provided by MOAB) as a means of doing faster spatial searches in order to determine matching tally bins. I was seeing very slow performance for building the k-d tree for large meshes and after digging into it determined that it has to do with some default parameters that MOAB uses for building the k-d tree. MOAB allows you to change these parameters but we don't expose this currently through our API. This PR adds an options attribute to UnstructuredMesh that will get passed to MOAB when building the k-d tree. So for example, you could specify:

mesh = openmc.UnstructuredMesh('mesh.h5m', 'moab', options='MAX_DEPTH=20')

which would tell MOAB to use a maximum tree depth of 20.

Further Details

I've found that there are two main parameters that impact performance, memory use, and initialization time for MOAB's k-d tree:

• PLANE_SET — this controls how planes are chosen to subdivide space in the k-d tree.
• MAX_DEPTH — as you might guess, this controls the maximum depth of the tree. Once the maximum depth is reached, all remaining points/elements are put together in a leaf node (which means a search would have to go through them sequentially).

There are four different options for PLANE_SET numbered 0–3 and the default is 1, in which the range of values is split by multiple planes that are evenly spaced and then the one that most evenly subdivides the range is chosen. In my testing, using PLANE_SET=2, in which a plane centered on the median vertex is used as the dividing plane, results in better performance, less memory use, and lower initialization time.

For MAX_DEPTH, the default value is 30 which seems to be overkill for all the cases I've tried. In this PR, we now use a maximum depth of 20 unless the user specifies otherwise.

As one particular example, when I run the model from our unstructured mesh example with fuel pins, I get the following performance characteristics:

You can see that with a maximum depth of 30, the performance is actually worse than for lower values. I've found 20 to be a sweet spot in terms of getting close to maximum performance while keeping initialization time and memory use modest.

Checklist

• I have performed a self-review of my own code
• I have run clang-format (version 15) 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)

[shimwell]

Interesting find Paul. Just wanted to ask if double down and embree makes a difference here.

[makeclean]

No, Embree has no K-D tree support (IIRC) - the K-D tree is built on tetrahedra while Embree operates on triangles. It would be some development effort but I expect you could hijack Embree BVH by inserting the triangular faces of the tetrahedra into the various calls, but there would be a lot of book keeping going from the triangle data you intersect back to the tets you actually want to record the data on.

[pshriwise]

No, Embree has no K-D tree support (IIRC) - the K-D tree is built on tetrahedra while Embree operates on triangles. It would be some development effort but I expect you could hijack Embree BVH by inserting the triangular faces of the tetrahedra into the various calls, but there would be a lot of book keeping going from the triangle data you intersect back to the tets you actually want to record the data on.

Spot on, @makeclean. This is how BVH's are used when ppl want to leverage RT cores on GPU for tetrahedral mesh.

[paulromano]

To be clear — I think we are in need of a better solution for accelerating searches on UM. This PR is just a bandaid to improve the status quo for the time being 🩹 I'm working with a UM for a stellarator that has > 1M elements and the k-d tree build times (and memory use) are just horrendous. The changes here at least get me to a point where I can run it in a reasonable amount of time.

[gridley]

Sounds like it'd be better to just track the rays on the unstructured mesh!

[gridley]

Looks good!

[pshriwise]

Looks good to me as well, @paulromano! I'm hoping to put a solution for larger meshes in place soon.