Random Ray Forward Flux Save in Adjoint Mode (#3962)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: jtramm

docs/source/usersguide/random_ray.rst

@@ -1105,10 +1105,15 @@ external source is present in the problem. Simulation settings (e.g., number of
 rays, batches, etc.) will be identical for both calculations. At the 
 conclusion of the run, all results (e.g., tallies, plots, etc.) will be 
 derived from the adjoint flux rather than the forward flux but are not labeled 
-any differently. The initial forward flux solution will not be stored or 
-available in the final statepoint file. Those wishing to do analysis requiring 
-both the forward and adjoint solutions will need to run two separate 
-simulations and load both statepoint files.
+any differently. When an initial forward solve is performed (i.e., when no
+user-specified adjoint source is present), its output files are also written to
+disk with a ``forward`` infix, so they are not overwritten by the subsequent
+adjoint solve. This applies to the statepoint, ``tallies.out``, and any voxel
+plots, e.g., ``statepoint.forward.N.h5`` and ``tallies.forward.out``; the
+adjoint solve keeps the usual file names. This allows analyses requiring both
+the forward and adjoint solutions to be performed from a single run. When
+generating FW-CADIS weight windows, no weight window file is written for the
+forward solve, as only the final adjoint-derived weight windows are meaningful.
 
 .. note::
     Use of the automated 

include/openmc/constants.h

@@ -368,6 +368,7 @@ enum class SolverType { MONTE_CARLO, RANDOM_RAY };
 enum class RandomRayVolumeEstimator { NAIVE, SIMULATION_AVERAGED, HYBRID };
 enum class RandomRaySourceShape { FLAT, LINEAR, LINEAR_XY };
 enum class RandomRaySampleMethod { PRNG, HALTON, S2 };
+enum class RandomRaySolve { FORWARD, FORWARD_FOR_ADJOINT, ADJOINT };
 
 //==============================================================================
 // Geometry Constants

include/openmc/random_ray/flat_source_domain.h

@@ -76,7 +76,10 @@ class FlatSourceDomain {
   //----------------------------------------------------------------------------
   // Static Data members
   static bool volume_normalized_flux_tallies_;
-  static bool adjoint_; // If the user wants outputs based on the adjoint flux
+  // If the user wants outputs based on the adjoint flux
+  static bool adjoint_requested_;
+  // The solve currently being executed
+  static RandomRaySolve solve_;
   static bool fw_cadis_local_;
   static double
     diagonal_stabilization_rho_; // Adjusts strength of diagonal stabilization

include/openmc/random_ray/random_ray_simulation.h

@@ -22,7 +22,7 @@ class RandomRaySimulation {
   void apply_fixed_sources_and_mesh_domains();
   void prepare_fw_fixed_sources_adjoint();
   void prepare_local_fixed_sources_adjoint();
-  void prepare_adjoint_simulation(bool fw_adjoint);
+  void prepare_adjoint_simulation(bool from_forward);
   void simulate();
   void output_simulation_results() const;
   void instability_check(

src/output.cpp

@@ -621,8 +621,15 @@ void write_tallies()
   if (model::tallies.empty())
     return;
 
+  // Tag tallies.out written during the forward solve of an adjoint run
+  const char* forward =
+    (FlatSourceDomain::solve_ == RandomRaySolve::FORWARD_FOR_ADJOINT)
+      ? "forward."
+      : "";
+
   // Set filename for tallies_out
-  std::string filename = fmt::format("{}tallies.out", settings::path_output);
+  std::string filename =
+    fmt::format("{}tallies.{}out", settings::path_output, forward);
 
   // Open the tallies.out file.
   std::ofstream tallies_out;

src/random_ray/flat_source_domain.cpp

@@ -30,7 +30,8 @@ namespace openmc {
 RandomRayVolumeEstimator FlatSourceDomain::volume_estimator_ {
   RandomRayVolumeEstimator::HYBRID};
 bool FlatSourceDomain::volume_normalized_flux_tallies_ {false};
-bool FlatSourceDomain::adjoint_ {false};
+bool FlatSourceDomain::adjoint_requested_ {false};
+RandomRaySolve FlatSourceDomain::solve_ {RandomRaySolve::FORWARD};
 bool FlatSourceDomain::fw_cadis_local_ {false};
 double FlatSourceDomain::diagonal_stabilization_rho_ {1.0};
 std::unordered_map<int, vector<std::pair<Source::DomainType, int>>>
@@ -556,7 +557,7 @@ double FlatSourceDomain::compute_fixed_source_normalization_factor() const
   // If we are in adjoint mode of a fixed source problem, the external
   // source is already normalized, such that all resulting fluxes are
   // also normalized.
-  if (adjoint_) {
+  if (solve_ == RandomRaySolve::ADJOINT) {
     return 1.0;
   }
 
@@ -795,6 +796,12 @@ void FlatSourceDomain::output_to_vtk() const
     double z_delta = width.z / Nz;
     std::string filename = openmc_plot->path_plot();
 
+    // Tag plots written during the forward solve of an adjoint run
+    if (solve_ == RandomRaySolve::FORWARD_FOR_ADJOINT) {
+      auto dot = filename.find_last_of('.');
+      filename = filename.substr(0, dot) + ".forward" + filename.substr(dot);
+    }
+
     // Perform sanity checks on file size
     uint64_t bytes = Nx * Ny * Nz * (negroups_ + 1 + 1 + 1) * sizeof(float);
     write_message(5, "Processing plot {}: {}... (Estimated size is {} MB)",
@@ -1002,9 +1009,10 @@ void FlatSourceDomain::output_to_vtk() const
 void FlatSourceDomain::apply_external_source_to_source_region(
   int src_idx, SourceRegionHandle& srh)
 {
-  auto s = (adjoint_ && !model::adjoint_sources.empty())
-             ? model::adjoint_sources[src_idx].get()
-             : model::external_sources[src_idx].get();
+  auto s =
+    (solve_ == RandomRaySolve::ADJOINT && !model::adjoint_sources.empty())
+      ? model::adjoint_sources[src_idx].get()
+      : model::external_sources[src_idx].get();
   auto is = dynamic_cast<IndependentSource*>(s);
   auto discrete = dynamic_cast<Discrete*>(is->energy());
   double strength_factor = is->strength();

src/random_ray/random_ray_simulation.cpp

@@ -189,7 +189,7 @@ void validate_random_ray_inputs()
 
   // Validate adjoint sources
   ///////////////////////////////////////////////////////////////////
-  if (FlatSourceDomain::adjoint_ && !model::adjoint_sources.empty()) {
+  if (FlatSourceDomain::adjoint_requested_ && !model::adjoint_sources.empty()) {
     for (int i = 0; i < model::adjoint_sources.size(); i++) {
       Source* s = model::adjoint_sources[i].get();
 
@@ -289,7 +289,8 @@ void openmc_finalize_random_ray()
 {
   FlatSourceDomain::volume_estimator_ = RandomRayVolumeEstimator::HYBRID;
   FlatSourceDomain::volume_normalized_flux_tallies_ = false;
-  FlatSourceDomain::adjoint_ = false;
+  FlatSourceDomain::adjoint_requested_ = false;
+  FlatSourceDomain::solve_ = RandomRaySolve::FORWARD;
   FlatSourceDomain::fw_cadis_local_ = false;
   FlatSourceDomain::fw_cadis_local_targets_.clear();
   FlatSourceDomain::mesh_domain_map_.clear();
@@ -356,19 +357,17 @@ void RandomRaySimulation::prepare_local_fixed_sources_adjoint()
   }
 }
 
-void RandomRaySimulation::prepare_adjoint_simulation(bool fw_adjoint)
+void RandomRaySimulation::prepare_adjoint_simulation(bool from_forward)
 {
   reset_timers();
 
   if (mpi::master)
     header("ADJOINT FLUX SOLVE", 3);
 
-  if (fw_adjoint) {
-    // Forward simulation has already been run;
-    // Configure the domain for adjoint simulation and
-    // re-initialize OpenMC general data structures
-    FlatSourceDomain::adjoint_ = true;
-
+  if (from_forward) {
+    // The forward solve has already run. Re-initialize OpenMC's general data
+    // structures for the adjoint solve and derive the adjoint source from the
+    // forward flux.
     openmc_simulation_init();
 
     prepare_fw_fixed_sources_adjoint();
@@ -603,7 +602,8 @@ void RandomRaySimulation::print_results_random_ray(
     }
     fmt::print(" Volume Estimator Type             = {}\n", estimator);
 
-    std::string adjoint_true = (FlatSourceDomain::adjoint_) ? "ON" : "OFF";
+    std::string adjoint_true =
+      (FlatSourceDomain::solve_ == RandomRaySolve::ADJOINT) ? "ON" : "OFF";
     fmt::print(" Adjoint Flux Mode                 = {}\n", adjoint_true);
 
     std::string shape;
@@ -675,60 +675,49 @@ void RandomRaySimulation::print_results_random_ray(
 
 void openmc_run_random_ray()
 {
-  //////////////////////////////////////////////////////////
-  // Run forward simulation
-  //////////////////////////////////////////////////////////
-
-  // Check if adjoint calculation is needed, and if local adjoint source(s)
-  // are present. If an adjoint calculation is needed and no sources are
-  // specified, we will run a forward calculation first to calculate adjoint
-  // sources for global variance reduction, then perform an adjoint
-  // calculation later.
-  bool adjoint_needed = openmc::FlatSourceDomain::adjoint_;
-  bool fw_adjoint = openmc::model::adjoint_sources.empty() && adjoint_needed;
-
-  // If we're going to do an adjoint simulation with forward-weighted adjoint
-  // sources afterwards, report that this is the initial forward flux solve.
-  if (!adjoint_needed || fw_adjoint) {
-    // Configure the domain for forward simulation
-    openmc::FlatSourceDomain::adjoint_ = false;
-
-    if (adjoint_needed && openmc::mpi::master)
-      openmc::header("FORWARD FLUX SOLVE", 3);
+  using namespace openmc;
+
+  // Determine which solves to run. If adjoint results are requested and no
+  // user-defined adjoint source is present, an initial forward solve is needed
+  // to construct the adjoint source from the forward flux (FW-CADIS). If the
+  // user has defined an adjoint source, the forward solve is skipped and only
+  // the adjoint solve is run.
+  const bool run_adjoint = FlatSourceDomain::adjoint_requested_;
+  const bool have_adjoint_source = !model::adjoint_sources.empty();
+  const bool run_forward = !(run_adjoint && have_adjoint_source);
+
+  // Set the initial solve type
+  if (!run_forward) {
+    FlatSourceDomain::solve_ = RandomRaySolve::ADJOINT;
+  } else if (run_adjoint) {
+    FlatSourceDomain::solve_ = RandomRaySolve::FORWARD_FOR_ADJOINT;
   } else {
-    // Configure domain for adjoint simulation (later)
-    openmc::FlatSourceDomain::adjoint_ = true;
+    FlatSourceDomain::solve_ = RandomRaySolve::FORWARD;
   }
 
   // Initialize OpenMC general data structures
   openmc_simulation_init();
 
   // Validate that inputs meet requirements for random ray mode
-  if (openmc::mpi::master)
-    openmc::validate_random_ray_inputs();
+  if (mpi::master)
+    validate_random_ray_inputs();
 
   // Initialize Random Ray Simulation Object
-  openmc::RandomRaySimulation sim;
+  RandomRaySimulation sim;
 
-  if (!adjoint_needed || fw_adjoint) {
-    // Initialize fixed sources, if present
+  // Run the forward solve
+  if (run_forward) {
+    // When an adjoint solve follows, report this as the initial forward solve
+    if (run_adjoint && mpi::master)
+      header("FORWARD FLUX SOLVE", 3);
     sim.apply_fixed_sources_and_mesh_domains();
-
-    // Execute random ray simulation
     sim.simulate();
   }
 
-  //////////////////////////////////////////////////////////
-  // Run adjoint simulation (if enabled)
-  //////////////////////////////////////////////////////////
-
-  if (!adjoint_needed) {
-    return;
+  // Run the adjoint solve
+  if (run_adjoint) {
+    FlatSourceDomain::solve_ = RandomRaySolve::ADJOINT;
+    sim.prepare_adjoint_simulation(run_forward);
+    sim.simulate();
   }
-
-  // Setup for adjoint simulation
-  sim.prepare_adjoint_simulation(fw_adjoint);
-
-  // Execute random ray simulation
-  sim.simulate();
 }

src/settings.cpp

@@ -322,7 +322,7 @@ void get_run_parameters(pugi::xml_node node_base)
         get_node_value_bool(random_ray_node, "volume_normalized_flux_tallies");
     }
     if (check_for_node(random_ray_node, "adjoint")) {
-      FlatSourceDomain::adjoint_ =
+      FlatSourceDomain::adjoint_requested_ =
         get_node_value_bool(random_ray_node, "adjoint");
     }
     if (check_for_node(random_ray_node, "sample_method")) {

src/simulation.cpp

@@ -16,6 +16,7 @@
 #include "openmc/particle.h"
 #include "openmc/photon.h"
 #include "openmc/random_lcg.h"
+#include "openmc/random_ray/flat_source_domain.h"
 #include "openmc/settings.h"
 #include "openmc/source.h"
 #include "openmc/state_point.h"
@@ -200,9 +201,12 @@ int openmc_simulation_finalize()
   if (settings::output_tallies && mpi::master)
     write_tallies();
 
-  // If weight window generators are present in this simulation,
-  // write a weight windows file
-  if (variance_reduction::weight_windows_generators.size() > 0) {
+  // If weight window generators are present in this simulation, write a
+  // weight windows file. This is skipped during the forward solve of an
+  // adjoint (FW-CADIS) run, where only the adjoint-derived weight windows
+  // are meaningful.
+  if (variance_reduction::weight_windows_generators.size() > 0 &&
+      FlatSourceDomain::solve_ != RandomRaySolve::FORWARD_FOR_ADJOINT) {
     openmc_weight_windows_export();
   }
 

src/state_point.cpp

@@ -22,6 +22,7 @@
 #include "openmc/nuclide.h"
 #include "openmc/output.h"
 #include "openmc/particle_type.h"
+#include "openmc/random_ray/flat_source_domain.h"
 #include "openmc/settings.h"
 #include "openmc/simulation.h"
 #include "openmc/tallies/derivative.h"
@@ -46,9 +47,15 @@ extern "C" int openmc_statepoint_write(const char* filename, bool* write_source)
     // Determine width for zero padding
     int w = std::to_string(settings::n_max_batches).size();
 
+    // Tag statepoints written during the forward solve of an adjoint run
+    const char* forward =
+      (FlatSourceDomain::solve_ == RandomRaySolve::FORWARD_FOR_ADJOINT)
+        ? "forward."
+        : "";
+
     // Set filename for state point
-    filename_ = fmt::format("{0}statepoint.{1:0{2}}.h5", settings::path_output,
-      simulation::current_batch, w);
+    filename_ = fmt::format("{0}statepoint.{3}{1:0{2}}.h5",
+      settings::path_output, simulation::current_batch, w, forward);
   }
 
   // If a file name was specified, ensure it has .h5 file extension