Parallelization of Weight Window Update (#3467)

This commit is contained in:
John Tramm 2025-06-25 08:46:40 -05:00 committed by GitHub
parent 2a0299aec7
commit 15dfe7e8b9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -6,11 +6,11 @@
#include <set>
#include <string>
#include "xtensor/xdynamic_view.hpp"
#include "xtensor/xindex_view.hpp"
#include "xtensor/xio.hpp"
#include "xtensor/xmasked_view.hpp"
#include "xtensor/xnoalias.hpp"
#include "xtensor/xstrided_view.hpp"
#include "xtensor/xview.hpp"
#include "openmc/error.h"
@ -506,8 +506,18 @@ void WeightWindows::update_weights(const Tally* tally, const std::string& value,
///////////////////////////
this->check_tally_update_compatibility(tally);
lower_ww_.fill(-1);
upper_ww_.fill(-1);
// Dimensions of weight window arrays
int e_bins = lower_ww_.shape()[0];
int64_t mesh_bins = lower_ww_.shape()[1];
// Initialize weight window arrays to -1.0 by default
#pragma omp parallel for collapse(2) schedule(static)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
lower_ww_(e, m) = -1.0;
upper_ww_(e, m) = -1.0;
}
}
// determine which value to use
const std::set<std::string> allowed_values = {"mean", "rel_err"};
@ -605,10 +615,12 @@ void WeightWindows::update_weights(const Tally* tally, const std::string& value,
}
// down-select data based on particle and score
auto sum = xt::view(transposed_view, particle_idx, xt::all(), xt::all(),
score_index, static_cast<int>(TallyResult::SUM));
auto sum_sq = xt::view(transposed_view, particle_idx, xt::all(), xt::all(),
score_index, static_cast<int>(TallyResult::SUM_SQ));
auto sum = xt::dynamic_view(
transposed_view, {particle_idx, xt::all(), xt::all(), score_index,
static_cast<int>(TallyResult::SUM)});
auto sum_sq = xt::dynamic_view(
transposed_view, {particle_idx, xt::all(), xt::all(), score_index,
static_cast<int>(TallyResult::SUM_SQ)});
int n = tally->n_realizations_;
//////////////////////////////////////////////
@ -626,78 +638,117 @@ void WeightWindows::update_weights(const Tally* tally, const std::string& value,
auto& new_bounds = this->lower_ww_;
auto& rel_err = this->upper_ww_;
// noalias avoids memory allocation here
xt::noalias(new_bounds) = sum / n;
xt::noalias(rel_err) =
xt::sqrt(((sum_sq / n) - xt::square(new_bounds)) / (n - 1)) / new_bounds;
xt::filter(rel_err, sum <= 0.0).fill(INFTY);
if (value == "rel_err")
xt::noalias(new_bounds) = 1 / rel_err;
// get mesh volumes
auto mesh_vols = this->mesh()->volumes();
int e_bins = new_bounds.shape()[0];
if (method == WeightWindowUpdateMethod::MAGIC) {
// If we are computing weight windows with forward fluxes derived from a
// Monte Carlo or forward random ray solve, we use the MAGIC algorithm.
for (int e = 0; e < e_bins; e++) {
// select all
auto group_view = xt::view(new_bounds, e);
// divide by volume of mesh elements
for (int i = 0; i < group_view.size(); i++) {
group_view[i] /= mesh_vols[i];
// Calculate mean (new_bounds) and relative error
#pragma omp parallel for collapse(2) schedule(static)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
// Calculate mean
new_bounds(e, m) = sum(e, m) / n;
// Calculate relative error
if (sum(e, m) > 0.0) {
double mean_val = new_bounds(e, m);
double variance = (sum_sq(e, m) / n - mean_val * mean_val) / (n - 1);
rel_err(e, m) = std::sqrt(variance) / mean_val;
} else {
rel_err(e, m) = INFTY;
}
double group_max =
*std::max_element(group_view.begin(), group_view.end());
// normalize values in this energy group by the maximum value for this
// group
if (group_max > 0.0)
group_view /= 2.0 * group_max;
}
} else {
// If we are computing weight windows with adjoint fluxes derived from an
// adjoint random ray solve, we use the FW-CADIS algorithm.
for (int e = 0; e < e_bins; e++) {
// select all
auto group_view = xt::view(new_bounds, e);
// divide by volume of mesh elements
for (int i = 0; i < group_view.size(); i++) {
group_view[i] /= mesh_vols[i];
if (value == "rel_err") {
new_bounds(e, m) = 1.0 / rel_err(e, m);
}
}
// We take the inverse, but are careful not to divide by zero e.g. if some
// mesh bins are not reachable in the physical geometry.
xt::noalias(new_bounds) =
xt::where(xt::not_equal(new_bounds, 0.0), 1.0 / new_bounds, 0.0);
auto max_val = xt::amax(new_bounds)();
xt::noalias(new_bounds) = new_bounds / (2.0 * max_val);
// For bins that were missed, we use the minimum weight window value. This
// shouldn't matter except for plotting.
auto min_val = xt::amin(new_bounds)();
xt::noalias(new_bounds) =
xt::where(xt::not_equal(new_bounds, 0.0), new_bounds, min_val);
}
// make sure that values where the mean is zero are set s.t. the weight window
// value will be ignored
xt::filter(new_bounds, sum <= 0.0).fill(-1.0);
// Divide by volume of mesh elements
#pragma omp parallel for collapse(2) schedule(static)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
new_bounds(e, m) /= mesh_vols[m];
}
}
// make sure the weight windows are ignored for any locations where the
// relative error is higher than the specified relative error threshold
xt::filter(new_bounds, rel_err > threshold).fill(-1.0);
if (method == WeightWindowUpdateMethod::MAGIC) {
// For MAGIC, weight windows are proportional to the forward fluxes.
// We normalize weight windows independently for each energy group.
// update the bounds of this weight window class
// noalias avoids additional memory allocation
xt::noalias(upper_ww_) = ratio * lower_ww_;
// Find group maximum and normalize (per energy group)
for (int e = 0; e < e_bins; e++) {
double group_max = 0.0;
// Find maximum value across all elements in this energy group
#pragma omp parallel for schedule(static) reduction(max : group_max)
for (int64_t m = 0; m < mesh_bins; m++) {
if (new_bounds(e, m) > group_max) {
group_max = new_bounds(e, m);
}
}
// Normalize values in this energy group by the maximum value
if (group_max > 0.0) {
double norm_factor = 1.0 / (2.0 * group_max);
#pragma omp parallel for schedule(static)
for (int64_t m = 0; m < mesh_bins; m++) {
new_bounds(e, m) *= norm_factor;
}
}
}
} else {
// For FW-CADIS, weight windows are inversely proportional to the adjoint
// fluxes. We normalize the weight windows across all energy groups.
#pragma omp parallel for collapse(2) schedule(static)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
// Take the inverse, but are careful not to divide by zero
if (new_bounds(e, m) != 0.0) {
new_bounds(e, m) = 1.0 / new_bounds(e, m);
} else {
new_bounds(e, m) = 0.0;
}
}
}
// Find the maximum value across all elements
double max_val = 0.0;
#pragma omp parallel for collapse(2) schedule(static) reduction(max : max_val)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
if (new_bounds(e, m) > max_val) {
max_val = new_bounds(e, m);
}
}
}
// Parallel normalization
if (max_val > 0.0) {
double norm_factor = 1.0 / (2.0 * max_val);
#pragma omp parallel for collapse(2) schedule(static)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
new_bounds(e, m) *= norm_factor;
}
}
}
}
// Final processing
#pragma omp parallel for collapse(2) schedule(static)
for (int e = 0; e < e_bins; e++) {
for (int64_t m = 0; m < mesh_bins; m++) {
// Values where the mean is zero should be ignored
if (sum(e, m) <= 0.0) {
new_bounds(e, m) = -1.0;
}
// Values where the relative error is higher than the threshold should be
// ignored
else if (rel_err(e, m) > threshold) {
new_bounds(e, m) = -1.0;
}
// Set the upper bounds
upper_ww_(e, m) = ratio * lower_ww_(e, m);
}
}
}
void WeightWindows::check_tally_update_compatibility(const Tally* tally)