diff --git a/include/openmc/geometry_aux.h b/include/openmc/geometry_aux.h index 28b19e315..899326a53 100644 --- a/include/openmc/geometry_aux.h +++ b/include/openmc/geometry_aux.h @@ -7,6 +7,7 @@ #include #include #include +#include #include "openmc/vector.h" @@ -63,9 +64,12 @@ int32_t find_root_universe(); //============================================================================== //! Populate all data structures needed for distribcells. +//! \param user_distribcells A set of cell indices to create distribcell data +//! structures for regardless of whether or not they are part of a tally +//! filter. //============================================================================== -void prepare_distribcell(); +void prepare_distribcell(const std::vector* user_distribcells = nullptr); //============================================================================== //! Recursively search through the geometry and count cell instances. diff --git a/include/openmc/tallies/filter_cell_instance.h b/include/openmc/tallies/filter_cell_instance.h index 789a4151e..387f7e591 100644 --- a/include/openmc/tallies/filter_cell_instance.h +++ b/include/openmc/tallies/filter_cell_instance.h @@ -44,6 +44,8 @@ public: const vector& cell_instances() const { return cell_instances_; } + const std::unordered_set& cells() const { return cells_; } + void set_cell_instances(gsl::span instances); private: diff --git a/src/geometry_aux.cpp b/src/geometry_aux.cpp index 9d3062901..857803421 100644 --- a/src/geometry_aux.cpp +++ b/src/geometry_aux.cpp @@ -311,12 +311,18 @@ find_root_universe() //============================================================================== void -prepare_distribcell() +prepare_distribcell(const std::vector* user_distribcells) { write_message("Preparing distributed cell instances...", 5); - // Find all cells listed in a DistribcellFilter or CellInstanceFilter std::unordered_set distribcells; + + // start with any cells manually specified via the C++ API + if (user_distribcells) { + distribcells.insert(user_distribcells->begin(), user_distribcells->end()); + } + + // Find all cells listed in a DistribcellFilter or CellInstanceFilter for (auto& filt : model::tally_filters) { auto* distrib_filt = dynamic_cast(filt.get()); auto* cell_inst_filt = dynamic_cast(filt.get()); @@ -324,7 +330,8 @@ prepare_distribcell() distribcells.insert(distrib_filt->cell()); } if (cell_inst_filt) { - for (const auto& c_inst : cell_inst_filt->cell_instances()) distribcells.insert(c_inst.index_cell); + const auto& filter_cells = cell_inst_filt->cells(); + distribcells.insert(filter_cells.begin(), filter_cells.end()); } }