Open-Sn · zhardy-lanl · Mar 4, 2024 · Feb 27, 2024
diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/acceleration/nl_keigen_acc_residual_func.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/acceleration/nl_keigen_acc_residual_func.cc
@@ -25,6 +25,7 @@ NLKEigenAccResidualFunction(SNES snes, Vec phi, Vec r, void* ctx)
   auto& q_moments_local = lbs_solver.QMomentsLocal();
   auto& phi_old_local = lbs_solver.PhiOldLocal();
   auto phi_temp = phi_old_local;
+  const auto& densities_local = lbs_solver.DensitiesLocal();
 
   const auto& phi_l = nl_context_ptr->phi_l_;
   const auto& phi_lph_i = nl_context_ptr->phi_lph_i_;
@@ -34,21 +35,24 @@ NLKEigenAccResidualFunction(SNES snes, Vec phi, Vec r, void* ctx)
 
   // Lambdas
   auto SetLBSFissionSource =
-    [&active_set_source_function, &front_gs](const VecDbl& input, VecDbl& output)
+    [&active_set_source_function, &front_gs, &densities_local](const VecDbl& input, VecDbl& output)
   {
     Set(output, 0.0);
-    active_set_source_function(
-      front_gs, output, input, APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
+    active_set_source_function(front_gs,
+                               output,
+                               input,
+                               densities_local,
+                               APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
   };
 
-  auto SetLBSScatterSource =
-    [&active_set_source_function, &front_gs](const VecDbl& input, VecDbl& output, bool suppress_wgs)
+  auto SetLBSScatterSource = [&active_set_source_function, &front_gs, &densities_local](
+                               const VecDbl& input, VecDbl& output, bool suppress_wgs)
   {
     Set(output, 0.0);
     SourceFlags source_flags = APPLY_AGS_SCATTER_SOURCES | APPLY_WGS_SCATTER_SOURCES;
     if (suppress_wgs)
       source_flags |= SUPPRESS_WG_SCATTER;
-    active_set_source_function(front_gs, output, input, source_flags);
+    active_set_source_function(front_gs, output, input, densities_local, source_flags);
   };
 
   auto SetPhi0FissionSource =

diff --git a/...es/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/nl_keigen_ags_residual_func.cc b/...es/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/nl_keigen_ags_residual_func.cc
@@ -21,6 +21,7 @@ NLKEigenResidualFunction(SNES snes, Vec phi, Vec r, void* ctx)
   auto& lbs_solver = nl_context_ptr->lbs_solver_;
   const auto& phi_old_local = lbs_solver.PhiOldLocal();
   auto& q_moments_local = lbs_solver.QMomentsLocal();
+  const auto& densities_local = lbs_solver.DensitiesLocal();
 
   auto active_set_source_function = lbs_solver.GetActiveSetSourceFunction();
 
@@ -38,6 +39,7 @@ NLKEigenResidualFunction(SNES snes, Vec phi, Vec r, void* ctx)
     active_set_source_function(groupset,
                                q_moments_local,
                                phi_old_local,
+                               densities_local,
                                lbs::APPLY_AGS_FISSION_SOURCES | lbs::APPLY_WGS_FISSION_SOURCES);
 
   const double k_eff = lbs_solver.ComputeFissionProduction(phi_old_local);
@@ -51,7 +53,8 @@ NLKEigenResidualFunction(SNES snes, Vec phi, Vec r, void* ctx)
     SourceFlags source_flags = lbs::APPLY_AGS_SCATTER_SOURCES | lbs::APPLY_WGS_SCATTER_SOURCES;
     if (supress_wgs)
       source_flags |= lbs::SUPPRESS_WG_SCATTER;
-    active_set_source_function(groupset, q_moments_local, phi_old_local, source_flags);
+    active_set_source_function(
+      groupset, q_moments_local, phi_old_local, densities_local, source_flags);
   }
 
   // Sweep all the groupsets

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/power_iteration_keigen.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/power_iteration_keigen.cc
@@ -34,6 +34,7 @@ PowerIterationKEigen(LBSSolver& lbs_solver, double tolerance, int max_iterations
   auto& q_moments_local = lbs_solver.QMomentsLocal();
   auto& phi_old_local = lbs_solver.PhiOldLocal();
   auto& phi_new_local = lbs_solver.PhiNewLocal();
+  const auto& densities_local = lbs_solver.DensitiesLocal();
   auto primary_ags_solver = lbs_solver.GetPrimaryAGSSolver();
   auto& groupsets = lbs_solver.Groupsets();
   auto active_set_source_function = lbs_solver.GetActiveSetSourceFunction();
@@ -59,6 +60,7 @@ PowerIterationKEigen(LBSSolver& lbs_solver, double tolerance, int max_iterations
       active_set_source_function(groupset,
                                  q_moments_local,
                                  phi_old_local,
+                                 densities_local,
                                  APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
 
     Scale(q_moments_local, 1.0 / k_eff);

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/power_iteration_keigen1.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/power_iteration_keigen1.cc
@@ -41,6 +41,7 @@ PowerIterationKEigen1(LBSSolver& lbs_solver, double tolerance, int max_iteration
   auto& q_moments_local = lbs_solver.QMomentsLocal();
   auto& phi_old_local = lbs_solver.PhiOldLocal();
   auto& phi_new_local = lbs_solver.PhiNewLocal();
+  const auto& densities_local = lbs_solver.DensitiesLocal();
   auto primary_ags_solver = lbs_solver.GetPrimaryAGSSolver();
   auto& groupsets = lbs_solver.Groupsets();
   auto active_set_source_function = lbs_solver.GetActiveSetSourceFunction();
@@ -67,26 +68,30 @@ PowerIterationKEigen1(LBSSolver& lbs_solver, double tolerance, int max_iteration
   double k_eff_change = 1.0;
 
   /**Lambda for the creation of fission sources.*/
-  auto SetLBSFissionSource = [&active_set_source_function, &front_gs, &q_moments_local](
-                               const VecDbl& input, const bool additive)
+  auto SetLBSFissionSource =
+    [&active_set_source_function, &front_gs, &q_moments_local, &densities_local](
+      const VecDbl& input, const bool additive)
   {
     if (not additive)
       Set(q_moments_local, 0.0);
-    active_set_source_function(
-      front_gs, q_moments_local, input, APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
+    active_set_source_function(front_gs,
+                               q_moments_local,
+                               input,
+                               densities_local,
+                               APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
   };
 
   /**Lambda for the creation of scattering sources*/
   auto SetLBSScatterSource =
-    [&active_set_source_function, &front_gs, &q_moments_local](
+    [&active_set_source_function, &front_gs, &q_moments_local, &densities_local](
       const VecDbl& input, const bool additive, const bool suppress_wgs = false)
   {
     if (not additive)
       Set(q_moments_local, 0.0);
     SourceFlags source_flags = APPLY_AGS_SCATTER_SOURCES | APPLY_WGS_SCATTER_SOURCES;
     if (suppress_wgs)
       source_flags |= SUPPRESS_WG_SCATTER;
-    active_set_source_function(front_gs, q_moments_local, input, source_flags);
+    active_set_source_function(front_gs, q_moments_local, input, densities_local, source_flags);
   };
 
   auto phi_temp = phi_old_local;

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/power_iteration_keigen2.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/power_iteration_keigen2.cc
@@ -43,6 +43,7 @@ PowerIterationKEigen2(LBSSolver& lbs_solver, double tolerance, int max_iteration
   auto& q_moments_local = lbs_solver.QMomentsLocal();
   auto& phi_old_local = lbs_solver.PhiOldLocal();
   auto& phi_new_local = lbs_solver.PhiNewLocal();
+  const auto& densities_local = lbs_solver.DensitiesLocal();
   auto primary_ags_solver = lbs_solver.GetPrimaryAGSSolver();
   auto& groupsets = lbs_solver.Groupsets();
   auto active_set_source_function = lbs_solver.GetActiveSetSourceFunction();
@@ -87,11 +88,14 @@ PowerIterationKEigen2(LBSSolver& lbs_solver, double tolerance, int max_iteration
 
   /**Lambda for the creation of fission sources.*/
   auto SetLBSFissionSource =
-    [&active_set_source_function, &front_gs](const VecDbl& input, VecDbl& output)
+    [&active_set_source_function, &front_gs, &densities_local](const VecDbl& input, VecDbl& output)
   {
     Set(output, 0.0);
-    active_set_source_function(
-      front_gs, output, input, APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
+    active_set_source_function(front_gs,
+                               output,
+                               input,
+                               densities_local,
+                               APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
   };
 
   // Start power iterations
@@ -121,6 +125,7 @@ PowerIterationKEigen2(LBSSolver& lbs_solver, double tolerance, int max_iteration
     active_set_source_function(front_gs,
                                q_moments_local,
                                phi_new_local,
+                               densities_local,
                                APPLY_AGS_SCATTER_SOURCES | APPLY_WGS_SCATTER_SOURCES);
 
     frons_wgs_context->ApplyInverseTransportOperator(SourceFlags());

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/wgs_context.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/wgs_context.cc
@@ -39,7 +39,11 @@ WGSContext::MatrixAction(Mat& matrix, Vec& action_vector, Vec& action)
   // Setting the source using updated phi_old
   auto& q_moments_local = lbs_solver_.QMomentsLocal();
   q_moments_local.assign(q_moments_local.size(), 0.0);
-  set_source_function_(groupset, q_moments_local, lbs_solver.PhiOldLocal(), lhs_src_scope_);
+  set_source_function_(groupset,
+                       q_moments_local,
+                       lbs_solver.PhiOldLocal(),
+                       lbs_solver.DensitiesLocal(),
+                       lhs_src_scope_);
 
   // Apply transport operator
   gs_context_ptr->ApplyInverseTransportOperator(lhs_src_scope_);

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/wgs_linear_solver.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/iterative_methods/wgs_linear_solver.cc
@@ -158,8 +158,11 @@ WGSLinearSolver::SetRHS()
   if (not single_richardson)
   {
     const auto scope = gs_context_ptr->rhs_src_scope_ | ZERO_INCOMING_DELAYED_PSI;
-    gs_context_ptr->set_source_function_(
-      groupset, lbs_solver.QMomentsLocal(), lbs_solver.PhiOldLocal(), scope);
+    gs_context_ptr->set_source_function_(groupset,
+                                         lbs_solver.QMomentsLocal(),
+                                         lbs_solver.PhiOldLocal(),
+                                         lbs_solver.DensitiesLocal(),
+                                         scope);
 
     // Apply transport operator
     gs_context_ptr->ApplyInverseTransportOperator(scope);
@@ -186,8 +189,11 @@ WGSLinearSolver::SetRHS()
   else
   {
     const auto scope = gs_context_ptr->rhs_src_scope_ | gs_context_ptr->lhs_src_scope_;
-    gs_context_ptr->set_source_function_(
-      groupset, lbs_solver.QMomentsLocal(), lbs_solver.PhiOldLocal(), scope);
+    gs_context_ptr->set_source_function_(groupset,
+                                         lbs_solver.QMomentsLocal(),
+                                         lbs_solver.PhiOldLocal(),
+                                         lbs_solver.DensitiesLocal(),
+                                         scope);
 
     // Apply transport operator
     gs_context_ptr->ApplyInverseTransportOperator(scope);

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/lbs_solver.cc b/modules/linear_boltzmann_solvers/a_lbs_solver/lbs_solver.cc
@@ -355,6 +355,18 @@ LBSSolver::PsiNewLocal() const
   return psi_new_local_;
 }
 
+std::vector<double>&
+LBSSolver::DensitiesLocal()
+{
+  return densities_local_;
+}
+
+const std::vector<double>&
+LBSSolver::DensitiesLocal() const
+{
+  return densities_local_;
+}
+
 const std::map<uint64_t, std::shared_ptr<SweepBndry>>&
 LBSSolver::SweepBoundaries() const
 {
@@ -872,6 +884,9 @@ LBSSolver::PerformInputChecks()
     options_.geometry_type = GeometryType::THREED_CARTESIAN;
   else
     OpenSnLogicalError("Cannot deduce geometry type from mesh.");
+
+  // Assign placeholder unit densities
+  densities_local_.assign(grid_ptr_->local_cells.size(), 1.0);
 }
 
 void
@@ -2423,7 +2438,8 @@ LBSSolver::MakeSourceMomentsFromPhi()
   {
     active_set_source_function_(groupset,
                                 source_moments,
-                                PhiOldLocal(),
+                                phi_old_local_,
+                                densities_local_,
                                 APPLY_AGS_SCATTER_SOURCES | APPLY_WGS_SCATTER_SOURCES |
                                   APPLY_AGS_FISSION_SOURCES | APPLY_WGS_FISSION_SOURCES);
   }

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/lbs_solver.h b/modules/linear_boltzmann_solvers/a_lbs_solver/lbs_solver.h
@@ -252,6 +252,16 @@ class LBSSolver : public opensn::Solver
    */
   const std::vector<VecDbl>& PsiNewLocal() const;
 
+  /**
+   * Read/write access to the cell-wise densities.
+   */
+  std::vector<double>& DensitiesLocal();
+
+  /**
+   * Read access to the cell-wise densities.
+   */
+  const std::vector<double>& DensitiesLocal() const;
+
   /**
    * Returns the sweep boundaries as a read only reference
    */
@@ -558,6 +568,7 @@ class LBSSolver : public opensn::Solver
   std::vector<double> phi_new_local_, phi_old_local_;
   std::vector<std::vector<double>> psi_new_local_;
   std::vector<double> precursor_new_local_;
+  std::vector<double> densities_local_;
 
   SetSourceFunction active_set_source_function_;
 

diff --git a/modules/linear_boltzmann_solvers/a_lbs_solver/lbs_structs.h b/modules/linear_boltzmann_solvers/a_lbs_solver/lbs_structs.h
@@ -196,10 +196,11 @@ enum class PhiSTLOption
 };
 
 class LBSGroupset;
-typedef std::function<void(LBSGroupset& groupset,
-                           std::vector<double>& destination_q,
+typedef std::function<void(const LBSGroupset& groupset,
+                           std::vector<double>& q,
                            const std::vector<double>& phi,
-                           SourceFlags source_flags)>
+                           const std::vector<double>& densities,
+                           const SourceFlags source_flags)>
   SetSourceFunction;
 
 class AGSSchemeEntry;