From 979c1648015f66329a0e517e1f8724ba40b9aef2 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Mon, 23 Oct 2023 13:20:05 -0500
Subject: [PATCH] Check for invalid domain IDs in volume calculations (#2742)

---
 src/volume_calc.cpp             | 33 ++++++++++++++++++++++++++++++++-
 tests/unit_tests/test_volume.py | 16 ++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/src/volume_calc.cpp b/src/volume_calc.cpp
index 8949d5c20d..8cd697a9d8 100644
--- a/src/volume_calc.cpp
+++ b/src/volume_calc.cpp
@@ -97,6 +97,31 @@ VolumeCalculation::VolumeCalculation(pugi::xml_node node)
 
 vector<VolumeCalculation::Result> VolumeCalculation::execute() const
 {
+  // Check to make sure domain IDs are valid
+  for (auto uid : domain_ids_) {
+    switch (domain_type_) {
+    case TallyDomain::CELL:
+      if (model::cell_map.find(uid) == model::cell_map.end()) {
+        throw std::runtime_error {fmt::format(
+          "Cell {} in volume calculation does not exist in geometry.", uid)};
+      }
+      break;
+    case TallyDomain::MATERIAL:
+      if (model::material_map.find(uid) == model::material_map.end()) {
+        throw std::runtime_error {fmt::format(
+          "Material {} in volume calculation does not exist in geometry.",
+          uid)};
+      }
+      break;
+    case TallyDomain::UNIVERSE:
+      if (model::universe_map.find(uid) == model::universe_map.end()) {
+        throw std::runtime_error {fmt::format(
+          "Universe {} in volume calculation does not exist in geometry.",
+          uid)};
+      }
+    }
+  }
+
   // Shared data that is collected from all threads
   int n = domain_ids_.size();
   vector<vector<uint64_t>> master_indices(
@@ -519,7 +544,13 @@ int openmc_calculate_volumes()
 
     // Run volume calculation
     const auto& vol_calc {model::volume_calcs[i]};
-    auto results = vol_calc.execute();
+    std::vector<VolumeCalculation::Result> results;
+    try {
+      results = vol_calc.execute();
+    } catch (const std::exception& e) {
+      set_errmsg(e.what());
+      return OPENMC_E_UNASSIGNED;
+    }
 
     if (mpi::master) {
       std::string domain_type;
diff --git a/tests/unit_tests/test_volume.py b/tests/unit_tests/test_volume.py
index f49bb35018..f9343a0b1a 100644
--- a/tests/unit_tests/test_volume.py
+++ b/tests/unit_tests/test_volume.py
@@ -13,3 +13,19 @@ def test_infinity_handling():
 
     with pytest.raises(ValueError, match="must be finite"):
         openmc.VolumeCalculation([cell1], 100, lower_left, upper_right)
+
+
+@pytest.mark.parametrize('cls', [openmc.Cell, openmc.Material, openmc.Universe])
+def test_invalid_id(run_in_tmpdir, cls):
+    m = openmc.Material()
+    m.add_nuclide('U235', 0.02)
+    sph = openmc.Sphere(boundary_type='vacuum')
+    cell = openmc.Cell(fill=m, region=-sph)
+    model = openmc.Model(geometry=openmc.Geometry([cell]))
+
+    # Apply volume calculation with unused domains
+    model.settings.volume_calculations = openmc.VolumeCalculation(
+        [cls()], 10000, *model.geometry.bounding_box)
+
+    with pytest.raises(RuntimeError):
+        model.calculate_volumes()