Merge pull request #196 from lanl/mg/refactoring_coolingfunction

refactoring LightBulb_functions
lanl · Jan 17, 2024 · 532be0f · 532be0f
2 parents 7b56701 + 867a302
commit 532be0f
Show file tree

Hide file tree

Showing 5 changed files with 154 additions and 128 deletions.
diff --git a/src/phoebus_driver.cpp b/src/phoebus_driver.cpp
@@ -698,11 +698,12 @@ TaskListStatus PhoebusDriver::RadiationPostStep() {
     if (do_lightbulb) {
       pdo_gain_reducer = rad->MutableParam<parthenon::AllReduce<bool>>("do_gain_reducer");
     }
-    TaskRegion &async_region = tc.AddRegion(num_independent_task_lists);
-    for (int ib = 0; ib < num_independent_task_lists; ib++) {
-      auto pmb = blocks[ib].get();
-      auto &tl = async_region[ib];
-      auto &sc0 = pmb->meshblock_data.Get(stage_name[integrator->nstages]);
+    // creating a new sync region for light bulb functions
+    const int num_partitions = pmesh->DefaultNumPartitions();
+    TaskRegion &sync_region_lb = tc.AddRegion(num_partitions);
+    for (int ib = 0; ib < num_partitions; ++ib) {
+      auto &sc0 = pmesh->mesh_data.GetOrAdd("base", ib);
+      auto &tl = sync_region_lb[ib];
       auto finish_gain_reducer = none;
       if (do_lightbulb) {
         auto calc_tau = tl.AddTask(none, radiation::LightBulbCalcTau, sc0.get());
@@ -717,9 +718,8 @@ TaskListStatus PhoebusDriver::RadiationPostStep() {
             tl.AddTask(start_gain_reducer, &parthenon::AllReduce<bool>::CheckReduce,
                        pdo_gain_reducer);
         int reg_dep_id = 0;
-        async_region.AddRegionalDependencies(reg_dep_id++, ib, finish_gain_reducer);
+        sync_region_lb.AddRegionalDependencies(reg_dep_id++, ib, finish_gain_reducer);
       }
-
       auto calculate_four_force =
           tl.AddTask(finish_gain_reducer, radiation::CoolingFunctionCalculateFourForce,
                      sc0.get(), dt);
@@ -976,12 +976,12 @@ TaskListStatus PhoebusDriver::MonteCarloStep() {
 
   // Finalization calls
   {
-    TaskCollection tc;
-    TaskRegion &async_region1 = tc.AddRegion(num_task_lists_executed_independently);
-    for (int ib = 0; ib < num_task_lists_executed_independently; ib++) {
-      auto pmb = blocks[ib].get();
-      auto &tl = async_region1[ib];
-      auto &sc0 = pmb->meshblock_data.Get(stage_name[integrator->nstages]);
+    const int num_partitions = pmesh->DefaultNumPartitions();
+    TaskRegion &sync_region = tc.AddRegion(num_partitions);
+    for (int ib = 0; ib < num_partitions; ++ib) {
+      auto &base = pmesh->mesh_data.GetOrAdd("base", ib);
+      auto &sc0 = pmesh->mesh_data.GetOrAdd("base", ib);
+      auto &tl = sync_region[ib];
       auto apply_four_force =
           tl.AddTask(none, radiation::ApplyRadiationFourForce, sc0.get(), dt);
     }

diff --git a/src/phoebus_utils/relativity_utils.hpp b/src/phoebus_utils/relativity_utils.hpp
@@ -70,14 +70,22 @@ GetLorentzFactor(const Real vcon[Geometry::NDSPACE],
  *
  * RETURN - Lorentz factor in normal observer frame
  */
+template <typename CoordinateSystem_t>
 KOKKOS_INLINE_FUNCTION Real GetLorentzFactor(const Real vcon[Geometry::NDSPACE],
-                                             const Geometry::CoordSysMeshBlock &system,
-                                             CellLocation loc, const int k, const int j,
-                                             const int i) {
+                                             const CoordinateSystem_t &system,
+                                             CellLocation loc, const int b, const int k,
+                                             const int j, const int i) {
   Real gamma[Geometry::NDSPACE][Geometry::NDSPACE];
-  system.Metric(loc, k, j, i, gamma);
+  system.Metric(loc, b, k, j, i, gamma);
   return GetLorentzFactor(vcon, gamma);
 }
+template <typename CoordinateSystem_t>
+KOKKOS_INLINE_FUNCTION Real GetLorentzFactor(const Real vcon[Geometry::NDSPACE],
+                                             const CoordinateSystem_t &system,
+                                             CellLocation loc, const int k, const int j,
+                                             const int i) {
+  return GetLorentzFactor(vcon, system, loc, 0, k, j, i);
+}
 
 /*
  * Calculate the coordinate frame four-velocity from the normal observer three-velocity
@@ -90,20 +98,28 @@ KOKKOS_INLINE_FUNCTION Real GetLorentzFactor(const Real vcon[Geometry::NDSPACE],
  * PARAM[IN] - i - X1 index of meshblock cell
  * PARAM[OUT] - u - Coordinate frame contravariant four-velocity
  */
-KOKKOS_INLINE_FUNCTION void GetFourVelocity(const Real v[3],
-                                            const Geometry::CoordSysMeshBlock &system,
-                                            CellLocation loc, const int k, const int j,
-                                            const int i, Real u[Geometry::NDFULL]) {
+template <typename CoordinateSystem_t>
+KOKKOS_INLINE_FUNCTION void
+GetFourVelocity(const Real v[3], const CoordinateSystem_t &system, CellLocation loc,
+                const int b, const int k, const int j, const int i,
+                Real u[Geometry::NDFULL]) {
   Real beta[Geometry::NDSPACE];
-  Real W = GetLorentzFactor(v, system, loc, k, j, i);
-  Real alpha = system.Lapse(loc, k, j, i);
-  system.ContravariantShift(loc, k, j, i, beta);
+  Real W = GetLorentzFactor(v, system, loc, b, k, j, i);
+  Real alpha = system.Lapse(loc, b, k, j, i);
+  system.ContravariantShift(loc, b, k, j, i, beta);
   u[0] = robust::ratio(W, std::abs(alpha));
   for (int l = 1; l < Geometry::NDFULL; ++l) {
     u[l] = v[l - 1] - u[0] * beta[l - 1];
   }
 }
 
+template <typename CoordinateSystem_t>
+KOKKOS_INLINE_FUNCTION void
+GetFourVelocity(const Real v[3], const CoordinateSystem_t &system, CellLocation loc,
+                const int k, const int j, const int i, Real u[Geometry::NDFULL]) {
+  GetFourVelocity(v, system, loc, 0, k, j, i, u);
+}
+
 /*
  * Compute square of magnetic field from primitive velocity and magnetic field
  *