GetFourVelocity changed to get Mesh as well as MeshBlock

lanl · Jan 12, 2024 · 2520b98 · 2520b98
1 parent 9e25eb4
commit 2520b98
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 17 deletions.
diff --git a/src/phoebus_driver.cpp b/src/phoebus_driver.cpp
@@ -719,7 +719,7 @@ TaskListStatus PhoebusDriver::RadiationPostStep() {
             tl.AddTask(start_gain_reducer, &parthenon::AllReduce<bool>::CheckReduce,
                        pdo_gain_reducer);
         int reg_dep_id = 0;
-        sync_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,
@@ -977,12 +977,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_lb[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
@@ -90,20 +90,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
  *

diff --git a/src/radiation/cooling_function.cpp b/src/radiation/cooling_function.cpp
@@ -277,8 +277,7 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshData<Real> *rc, const double dt
           Tetrads.TetradToCoordCov(Gcov_tetrad, Gcov_coord);
           for (int mu = 0; mu <= 3; mu++) {
             // detg included above
-            Kokkos::atomic_add(&(v(b, iv::Gcov(mu), k, j, i)),
-                               -Gcov_coord[iv::Gcov(mu) - iv::Gcov(0)]);
+            Kokkos::atomic_add(&(v(b, iv::Gcov(mu), k, j, i)), -Gcov_coord[iv::Gcov(mu)]);
           }
           v(b, iv::GcovHeat(), k, j, i) =
               v(b, p::density(), k, j, i) * density_conversion_factor * heat;
@@ -322,7 +321,7 @@ TaskStatus CoolingFunctionCalculateFourForce(MeshData<Real> *rc, const double dt
 
               for (int mu = 0; mu <= 3; mu++) {
                 Kokkos::atomic_add(&(v(b, iv::Gcov(mu), k, j, i)),
-                                   -detG * Gcov_coord[iv::Gcov(mu) - iv::Gcov(0)]);
+                                   -detG * Gcov_coord[iv::Gcov(mu)]);
               }
               Kokkos::atomic_add(&(v(b, iv::Gye(), k, j, i)),
                                  -LeptonSign(s) * detG * Jye);