Adding runtime xs-map swapping for time-dependent problems

modules/linear_boltzmann_solvers/lbs_problem/lbs_problem.cc

@@ -285,6 +285,15 @@ LBSProblem::GetBlockID2XSMap() const
   return block_id_to_xs_map_;
 }
 
+void
+LBSProblem::SetBlockID2XSMap(const BlockID2XSMap& xs_map)
+{
+  block_id_to_xs_map_ = xs_map;
+  InitializeMaterials();
+  ResetGPUCarriers();
+  InitializeGPUExtras();
+}
+
 std::shared_ptr<MeshContinuum>
 LBSProblem::GetGrid() const
 {

modules/linear_boltzmann_solvers/lbs_problem/lbs_problem.h

@@ -136,6 +136,9 @@ class LBSProblem : public Problem
   /// Returns a reference to the map of material ids to XSs.
   const BlockID2XSMap& GetBlockID2XSMap() const;
 
+  /// Replaces the map of block ids to XSs and refreshes material data.
+  void SetBlockID2XSMap(const BlockID2XSMap& xs_map);
+
   /// Obtains a reference to the grid.
   std::shared_ptr<MeshContinuum> GetGrid() const;
 

python/lib/solver.cc

@@ -463,6 +463,46 @@ WrapLBS(py::module& slv)
         List of new volumetric sources to be added to the problem.
     )"
   );
+  lbs_problem.def(
+    "SetXSMap",
+    [](LBSProblem& self, py::kwargs& params)
+    {
+      BlockID2XSMap xs_map;
+      for (auto [key, value] : params)
+      {
+        auto c_key = key.cast<std::string>();
+        if (c_key == "xs_map")
+        {
+          auto xs_entries = value.cast<py::list>();
+          for (auto entry : xs_entries)
+          {
+            InputParameters xs_entry_pars = LBSProblem::GetXSMapEntryBlock();
+            xs_entry_pars.AssignParameters(pyobj_to_param_block("", entry.cast<py::dict>()));
+            const auto& block_ids =
+              xs_entry_pars.GetParam("block_ids").GetVectorValue<unsigned int>();
+            auto xs = xs_entry_pars.GetSharedPtrParam<MultiGroupXS>("xs");
+            for (const auto& block_id : block_ids)
+              xs_map[block_id] = xs;
+          }
+        }
+        else
+          throw std::runtime_error("Invalid argument provided to SetXSMap.\n");
+      }
+      self.SetBlockID2XSMap(xs_map);
+    },
+    R"(
+    Replace the block-id to cross-section map.
+
+    Parameters
+    ----------
+    xs_map: List[Dict]
+        A list of block-id to cross-section mapping dictionaries. Each dictionary supports:
+          - block_ids: List[int] (required)
+              Mesh block ids to associate with the cross section.
+          - xs: pyopensn.xs.MultiGroupXS (required)
+              Cross section object.
+    )"
+  );
   lbs_problem.def(
     "SetAdjoint",
     [](LBSProblem& self, bool adjoint)

test/python/modules/linear_boltzmann_solvers/transport_td/simple_2g_downscatter_td_swap.cxs

@@ -0,0 +1,24 @@
+NUM_GROUPS 2
+NUM_MOMENTS 1
+
+SIGMA_T_BEGIN
+0 2.0    # group 0
+1 1.2    # group 1
+SIGMA_T_END
+
+VELOCITY_BEGIN
+0 2.0    # v0
+1 0.5    # v1
+VELOCITY_END
+
+TRANSFER_MOMENTS_BEGIN
+# M_GFROM_GTO_VAL ell gfrom gto value  (ell = 0 only)
+# group 0 -> 0: no self-scatter
+M_GFROM_GTO_VAL   0   0   0   0.0
+# group 0 -> 1: downscatter
+M_GFROM_GTO_VAL   0   0   1   0.6
+# group 1 -> 0: no upscatter
+M_GFROM_GTO_VAL   0   1   0   0.0
+# group 1 -> 1: no self-scatter
+M_GFROM_GTO_VAL   0   1   1   0.0
+TRANSFER_MOMENTS_END

test/python/modules/linear_boltzmann_solvers/transport_td/tests.json

@@ -38,6 +38,19 @@
       }
     ]
   },
+  {
+    "file": "transport_3d_1g_v1_inf_med_td_swap.py",
+    "comment": "3D unstructured mesh, 1 group, vel = 1.0, infinte medium, time dependent, swap xs at 0.5s",
+    "num_procs": 4,
+    "checks": [
+      {
+        "type": "KeyValuePair",
+        "key": "Max phi(1s) = ",
+        "goldvalue": 1.736077,
+        "abs_tol": 1e-06
+      }
+    ]
+  },
   {
     "file": "transport_3d_1g_v2_inf_med_td.py",
     "comment": "3D unstructured mesh, 1 group, vel = 2.0, infinte medium, time dependent",
@@ -108,6 +121,25 @@
       }
     ]
   },
+  {
+    "file": "transport_3d_2g_inf_med_downscatter_td_swap.py",
+    "comment": "3D unstructured mesh, 2 group, downscatter, infinte medium, time dependent, swap xs at 0.5s",
+    "num_procs": 4,
+    "checks": [
+      {
+        "type": "KeyValuePair",
+        "key": "Max phi0(1s) = ",
+        "goldvalue": 1.978816,
+        "abs_tol": 1e-06
+      },
+      {
+        "type": "KeyValuePair",
+        "key": "Max phi1(1s) = ",
+        "goldvalue": 0.394171,
+        "abs_tol": 1e-06
+      }
+    ]
+  },
   {
     "file": "transport_3d_openmc_xs.py",
     "comment": "3D orthogonal mesh, 30 group, time dependent, OpenMC cross sections",