From 40f55a4c5f9837a43fcb4674aaf827ffec381f12 Mon Sep 17 00:00:00 2001
From: Pierre Kestener <pierre.kestener@cea.fr>
Date: Wed, 1 May 2024 21:09:28 +0200
Subject: [PATCH] Update MDrangePolicy more consistently.

---
 src/muscl/HydroInitFunctors2D.h | 123 ++++++++----------
 src/muscl/HydroInitFunctors3D.h | 122 +++++++-----------
 src/muscl/MHDInitFunctors2D.h   | 155 ++++++++++-------------
 src/muscl/MHDInitFunctors3D.h   | 214 +++++++++++++++-----------------
 src/muscl/MHDRunFunctors2D.h    |  10 +-
 src/muscl/SolverHydroMuscl.cpp  |   4 +-
 src/muscl/SolverHydroMuscl.h    |  15 +--
 src/muscl/SolverMHDMuscl.h      |  21 ++--
 8 files changed, 280 insertions(+), 384 deletions(-)

diff --git a/src/muscl/HydroInitFunctors2D.h b/src/muscl/HydroInitFunctors2D.h
index 1480f57..198b8a7 100644
--- a/src/muscl/HydroInitFunctors2D.h
+++ b/src/muscl/HydroInitFunctors2D.h
@@ -43,19 +43,20 @@ class InitImplodeFunctor2D : public HydroBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, ImplodeParams iparams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, ImplodeParams iparams, DataArray2d Udata)
   {
     InitImplodeFunctor2D functor(params, iparams, Udata);
-    Kokkos::parallel_for("InitImplodeFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitImplodeFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -77,9 +78,6 @@ class InitImplodeFunctor2D : public HydroBaseFunctor2D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -139,19 +137,20 @@ class InitBlastFunctor2D : public HydroBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, BlastParams bParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, BlastParams bParams, DataArray2d Udata)
   {
     InitBlastFunctor2D functor(params, bParams, Udata);
-    Kokkos::parallel_for("InitBlastFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitBlastFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -167,6 +166,7 @@ class InitBlastFunctor2D : public HydroBaseFunctor2D
 
     const real_t xmin = params.xmin;
     const real_t ymin = params.ymin;
+
     const real_t dx = params.dx;
     const real_t dy = params.dy;
 
@@ -182,9 +182,6 @@ class InitBlastFunctor2D : public HydroBaseFunctor2D
     const real_t blast_pressure_in = bParams.blast_pressure_in;
     const real_t blast_pressure_out = bParams.blast_pressure_out;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -228,20 +225,20 @@ class InitKelvinHelmholtzFunctor2D : public HydroBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, KHParams khParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, KHParams khParams, DataArray2d Udata)
   {
     InitKelvinHelmholtzFunctor2D functor(params, khParams, Udata);
     Kokkos::parallel_for(
-      "InitKelvinHelmholtzFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+      "InitKelvinHelmholtzFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -271,9 +268,6 @@ class InitKelvinHelmholtzFunctor2D : public HydroBaseFunctor2D
     const real_t ampl = khParams.amplitude;
     const real_t pressure = khParams.pressure;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -366,19 +360,20 @@ class InitGreshoVortexFunctor2D : public HydroBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, GreshoParams gvParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, GreshoParams gvParams, DataArray2d Udata)
   {
     InitGreshoVortexFunctor2D functor(params, gvParams, Udata);
-    Kokkos::parallel_for("InitGreshoVortexFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitGreshoVortexFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -406,9 +401,6 @@ class InitGreshoVortexFunctor2D : public HydroBaseFunctor2D
 
     const real_t p0 = rho0 / (gamma0 * Ma * Ma);
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -490,20 +482,20 @@ class InitFourQuadrantFunctor2D : public HydroBaseFunctor2D
         HydroState  U2,
         HydroState  U3,
         real_t      xt,
-        real_t      yt,
-        int         nbCells)
+        real_t      yt)
   {
     InitFourQuadrantFunctor2D functor(params, Udata, configNumber, U0, U1, U2, U3, xt, yt);
-    Kokkos::parallel_for("InitFourQuadrantFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitFourQuadrantFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -522,9 +514,6 @@ class InitFourQuadrantFunctor2D : public HydroBaseFunctor2D
     const real_t dx = params.dx;
     const real_t dy = params.dy;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -591,20 +580,20 @@ class InitIsentropicVortexFunctor2D : public HydroBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, IsentropicVortexParams iparams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, IsentropicVortexParams iparams, DataArray2d Udata)
   {
     InitIsentropicVortexFunctor2D functor(params, iparams, Udata);
     Kokkos::parallel_for(
-      "InitIsentropicVortexFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+      "InitIsentropicVortexFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -625,9 +614,6 @@ class InitIsentropicVortexFunctor2D : public HydroBaseFunctor2D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -699,19 +685,18 @@ class RayleighTaylorInstabilityFunctor2D : public HydroBaseFunctor2D
         DataArray2d                     Udata,
         VectorField2d                   gravity)
   {
-    uint64_t                           nbCells = params.isize * params.jsize;
     RayleighTaylorInstabilityFunctor2D functor(params, rtiparams, Udata, gravity);
     Kokkos::parallel_for(
-      "RayleighTaylorInstabilityFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+      "RayleighTaylorInstabilityFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -739,9 +724,6 @@ class RayleighTaylorInstabilityFunctor2D : public HydroBaseFunctor2D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -827,18 +809,18 @@ class RisingBubbleFunctor2D : public HydroBaseFunctor2D
   static void
   apply(HydroParams params, RisingBubbleParams rbparams, DataArray2d Udata, VectorField2d gravity)
   {
-    uint64_t              nbCells = params.isize * params.jsize;
     RisingBubbleFunctor2D functor(params, rbparams, Udata, gravity);
-    Kokkos::parallel_for("RisingBubbleFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "RisingBubbleFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -863,9 +845,6 @@ class RisingBubbleFunctor2D : public HydroBaseFunctor2D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -948,18 +927,19 @@ class InitDiskFunctor2D : public HydroBaseFunctor2D
         DataArray2d        Udata,
         VectorField2d      gravity)
   {
-    uint64_t          nbCells = params.isize * params.jsize;
     InitDiskFunctor2D functor(params, dparams, grav, Udata, gravity);
-    Kokkos::parallel_for("InitDiskFunctor2D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitDiskFunctor2D",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -991,9 +971,6 @@ class InitDiskFunctor2D : public HydroBaseFunctor2D
     const real_t rho_contrast = dparams.contrast_density;
     const real_t contrast_width = dparams.contrast_width;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx - xc;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy - yc;
     real_t z = 0;
diff --git a/src/muscl/HydroInitFunctors3D.h b/src/muscl/HydroInitFunctors3D.h
index 8dfe18c..369e1c1 100644
--- a/src/muscl/HydroInitFunctors3D.h
+++ b/src/muscl/HydroInitFunctors3D.h
@@ -37,25 +37,20 @@ class InitFakeFunctor3D : public HydroBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, DataArray3d Udata)
   {
     InitFakeFunctor3D functor(params, Udata);
-    Kokkos::parallel_for("InitFakeFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("InitFakeFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
-
-    int i, j, k;
-
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     Udata(i, j, k, ID) = 0.0;
     Udata(i, j, k, IP) = 0.0;
     Udata(i, j, k, IU) = 0.0;
@@ -82,20 +77,20 @@ class InitImplodeFunctor3D : public HydroBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, ImplodeParams iparams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, ImplodeParams iparams, DataArray3d Udata)
   {
     InitImplodeFunctor3D functor(params, iparams, Udata);
-    Kokkos::parallel_for("InitImplodeFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("InitImplodeFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -123,9 +118,6 @@ class InitImplodeFunctor3D : public HydroBaseFunctor3D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -190,20 +182,20 @@ class InitBlastFunctor3D : public HydroBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, BlastParams bParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, BlastParams bParams, DataArray3d Udata)
   {
     InitBlastFunctor3D functor(params, bParams, Udata);
-    Kokkos::parallel_for("InitBlastFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("InitBlastFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -241,10 +233,6 @@ class InitBlastFunctor3D : public HydroBaseFunctor3D
     const real_t blast_pressure_in = bParams.blast_pressure_in;
     const real_t blast_pressure_out = bParams.blast_pressure_out;
 
-
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -292,21 +280,20 @@ class InitKelvinHelmholtzFunctor3D : public HydroBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, KHParams khParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, KHParams khParams, DataArray3d Udata)
   {
     InitKelvinHelmholtzFunctor3D functor(params, khParams, Udata);
-    Kokkos::parallel_for(
-      "InitKelvinHelmholtzFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("InitKelvinHelmholtzFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -344,9 +331,6 @@ class InitKelvinHelmholtzFunctor3D : public HydroBaseFunctor3D
     const real_t ampl = khParams.amplitude;
     const real_t pressure = khParams.pressure;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -453,20 +437,20 @@ class InitGreshoVortexFunctor3D : public HydroBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, GreshoParams gvParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, GreshoParams gvParams, DataArray3d Udata)
   {
     InitGreshoVortexFunctor3D functor(params, gvParams, Udata);
-    Kokkos::parallel_for("InitGreshoVortexFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("InitGreshoVortexFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -502,9 +486,6 @@ class InitGreshoVortexFunctor3D : public HydroBaseFunctor3D
 
     const real_t p0 = rho0 / (gamma0 * Ma * Ma);
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     // real_t z = zmin + dz/2 + (k+nz*k_mpi-ghostWidth)*dz;
@@ -584,20 +565,18 @@ class RayleighTaylorInstabilityFunctor3D : public HydroBaseFunctor3D
         DataArray3d                     Udata,
         VectorField3d                   gravity)
   {
-    uint64_t                           nbCells = params.isize * params.jsize * params.ksize;
     RayleighTaylorInstabilityFunctor3D functor(params, rtiparams, Udata, gravity);
-    Kokkos::parallel_for(
-      "RayleighTaylorInstabilityFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("RayleighTaylorInstabilityFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -632,9 +611,6 @@ class RayleighTaylorInstabilityFunctor3D : public HydroBaseFunctor3D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -723,19 +699,18 @@ class RisingBubbleFunctor3D : public HydroBaseFunctor3D
   static void
   apply(HydroParams params, RisingBubbleParams rbparams, DataArray3d Udata, VectorField3d gravity)
   {
-    uint64_t              nbCells = params.isize * params.jsize * params.ksize;
     RisingBubbleFunctor3D functor(params, rbparams, Udata, gravity);
-    Kokkos::parallel_for("RisingBubbleFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("RisingBubbleFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -766,9 +741,6 @@ class RisingBubbleFunctor3D : public HydroBaseFunctor3D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -855,19 +827,18 @@ class InitDiskFunctor3D : public HydroBaseFunctor3D
         DataArray3d        Udata,
         VectorField3d      gravity)
   {
-    uint64_t          nbCells = params.isize * params.jsize * params.ksize;
     InitDiskFunctor3D functor(params, dparams, grav, Udata, gravity);
-    Kokkos::parallel_for("InitDiskFunctor3D", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for("InitDiskFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -905,9 +876,6 @@ class InitDiskFunctor3D : public HydroBaseFunctor3D
     const real_t rho_contrast = dparams.contrast_density;
     const real_t contrast_width = dparams.contrast_width;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx - xc;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy - yc;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz - zc;
diff --git a/src/muscl/MHDInitFunctors2D.h b/src/muscl/MHDInitFunctors2D.h
index fa185d6..09f38ce 100644
--- a/src/muscl/MHDInitFunctors2D.h
+++ b/src/muscl/MHDInitFunctors2D.h
@@ -44,19 +44,20 @@ class InitImplodeFunctor2D_MHD : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, ImplodeParams iparams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, ImplodeParams iparams, DataArray2d Udata)
   {
     InitImplodeFunctor2D_MHD functor(params, iparams, Udata);
-    Kokkos::parallel_for("InitImplodeFunctor2D_MHD", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitImplodeFunctor2D_MHD",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -78,9 +79,6 @@ class InitImplodeFunctor2D_MHD : public MHDBaseFunctor2D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -138,7 +136,7 @@ class InitImplodeFunctor2D_MHD : public MHDBaseFunctor2D
   ImplodeParams iparams;
   DataArray2d   Udata;
 
-}; // InitImplodeFunctor
+}; // InitImplodeFunctor2D_MHD
 
 /*************************************************/
 /*************************************************/
@@ -154,19 +152,20 @@ class InitBlastFunctor2D_MHD : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, BlastParams bParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, BlastParams bParams, DataArray2d Udata)
   {
     InitBlastFunctor2D_MHD functor(params, bParams, Udata);
-    Kokkos::parallel_for("InitBlastFunctor2D_MHD", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitBlastFunctor2D_MHD",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -198,10 +197,6 @@ class InitBlastFunctor2D_MHD : public MHDBaseFunctor2D
     const real_t blast_pressure_in = bParams.blast_pressure_in;
     const real_t blast_pressure_out = bParams.blast_pressure_out;
 
-
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -262,36 +257,38 @@ class InitOrszagTangFunctor2D : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, OrszagTangParams otParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, OrszagTangParams otParams, DataArray2d Udata)
   {
     InitOrszagTangFunctor2D functor(params, otParams, Udata);
 
     functor.phase = INIT_ALL_VAR_BUT_ENERGY;
-    Kokkos::parallel_for("InitOrszagTangFunctor2D INIT_ALL_VAR_BUT_ENERGY",
-                         Kokkos::RangePolicy<>(0, nbCells),
-                         functor);
+    Kokkos::parallel_for(
+      "InitOrszagTangFunctor2D INIT_ALL_VAR_BUT_ENERGY",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
 
     functor.phase = INIT_ENERGY;
-    Kokkos::parallel_for("InitOrszagTangFunctor2D INIT_ENERGY            ",
-                         Kokkos::RangePolicy<>(0, nbCells),
-                         functor);
+    Kokkos::parallel_for(
+      "InitOrszagTangFunctor2D INIT_ENERGY",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
     if (phase == INIT_ALL_VAR_BUT_ENERGY)
-      init_all_var_but_energy(index);
+      init_all_var_but_energy(i, j);
     else if (phase == INIT_ENERGY)
-      init_energy(index);
+      init_energy(i, j);
 
   } // end operator ()
 
   KOKKOS_INLINE_FUNCTION
   void
-  init_all_var_but_energy(const int index) const
+  init_all_var_but_energy(const int & i, const int & j) const
   {
 
     const int isize = params.isize;
@@ -322,9 +319,6 @@ class InitOrszagTangFunctor2D : public MHDBaseFunctor2D
     const double d0 = gamma0 * p0;
     const double v0 = 1.0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     double xPos = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     double yPos = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -353,7 +347,7 @@ class InitOrszagTangFunctor2D : public MHDBaseFunctor2D
 
   KOKKOS_INLINE_FUNCTION
   void
-  init_energy(const int index) const
+  init_energy(const int & i, const int & j) const
   {
 
     const int isize = params.isize;
@@ -374,9 +368,6 @@ class InitOrszagTangFunctor2D : public MHDBaseFunctor2D
     // const double d0    = gamma0*p0;
     // const double v0    = 1.0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     // double xPos = xmin + dx/2 + (i+nx*i_mpi-ghostWidth)*dx;
     // double yPos = ymin + dy/2 + (j+ny*j_mpi-ghostWidth)*dy;
 
@@ -413,20 +404,20 @@ class InitKelvinHelmholtzFunctor2D_MHD : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, KHParams khParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, KHParams khParams, DataArray2d Udata)
   {
     InitKelvinHelmholtzFunctor2D_MHD functor(params, khParams, Udata);
     Kokkos::parallel_for(
-      "InitKelvinHelmholtzFunctor2D_MHD", Kokkos::RangePolicy<>(0, nbCells), functor);
+      "InitKelvinHelmholtzFunctor2D_MHD",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -456,9 +447,6 @@ class InitKelvinHelmholtzFunctor2D_MHD : public MHDBaseFunctor2D
     const real_t ampl = khParams.amplitude;
     const real_t pressure = khParams.pressure;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -608,19 +596,20 @@ class InitRotorFunctor2D_MHD : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, RotorParams rParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, RotorParams rParams, DataArray2d Udata)
   {
     InitRotorFunctor2D_MHD functor(params, rParams, Udata);
-    Kokkos::parallel_for("InitRotorFunctor2D_MHD", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitRotorFunctor2D_MHD",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -654,9 +643,6 @@ class InitRotorFunctor2D_MHD : public MHDBaseFunctor2D
     const real_t xCenter = (xmax + xmin) / 2;
     const real_t yCenter = (ymax + ymin) / 2;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -745,10 +731,7 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
   struct TagInitEnergy
   {};
 
-  InitFieldLoopFunctor2D_MHD(HydroParams     params,
-                             FieldLoopParams flParams,
-                             DataArray2d     Udata,
-                             int             nbCells)
+  InitFieldLoopFunctor2D_MHD(HydroParams params, FieldLoopParams flParams, DataArray2d Udata)
     : MHDBaseFunctor2D(params)
     , flParams(flParams)
     , Udata(Udata)
@@ -758,52 +741,53 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, FieldLoopParams flParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, FieldLoopParams flParams, DataArray2d Udata)
   {
-    InitFieldLoopFunctor2D_MHD functor(params, flParams, Udata, nbCells);
+    InitFieldLoopFunctor2D_MHD functor(params, flParams, Udata);
 
     functor.phase = COMPUTE_VECTOR_POTENTIAL;
-    Kokkos::parallel_for("InitFieldLoopFunctor2D_MHD COMPUTE_VECTOR_POTENTIAL",
-                         Kokkos::RangePolicy<>(0, nbCells),
-                         functor);
+    Kokkos::parallel_for(
+      "InitFieldLoopFunctor2D_MHD COMPUTE_VECTOR_POTENTIAL",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
 
     functor.phase = DO_INIT_CONDITION;
-    Kokkos::parallel_for("InitFieldLoopFunctor2D_MHD DO_INIT_CONDITION       ",
-                         Kokkos::RangePolicy<>(0, nbCells),
-                         functor);
+    Kokkos::parallel_for(
+      "InitFieldLoopFunctor2D_MHD DO_INIT_CONDITION",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
 
     functor.phase = DO_INIT_ENERGY;
-    Kokkos::parallel_for("InitFieldLoopFunctor2D_MHD DO_INIT_ENERGY          ",
-                         Kokkos::RangePolicy<>(0, nbCells),
-                         functor);
+    Kokkos::parallel_for(
+      "InitFieldLoopFunctor2D_MHD DO_INIT_ENERGY",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
 
   } // apply
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {
     if (phase == COMPUTE_VECTOR_POTENTIAL)
     {
-      compute_vector_potential(index);
+      compute_vector_potential(i, j);
     }
     else if (phase == DO_INIT_CONDITION)
     {
-      do_init_condition(index);
+      do_init_condition(i, j);
     }
     else if (phase == DO_INIT_ENERGY)
     {
-      do_init_energy(index);
+      do_init_energy(i, j);
     }
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  compute_vector_potential(const int & index) const
+  compute_vector_potential(const int & i, const int & j) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
     const int ghostWidth = params.ghostWidth;
 
     const int nx = params.nx;
@@ -827,9 +811,6 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
     const real_t radius = flParams.radius;
     const real_t amplitude = flParams.amplitude;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -847,7 +828,7 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
 
   KOKKOS_INLINE_FUNCTION
   void
-  do_init_condition(const int & index) const
+  do_init_condition(const int & i, const int & j) const
   {
 
     const int isize = params.isize;
@@ -883,10 +864,6 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
     const real_t cos_theta = 2.0 / sqrt(5.0);
     const real_t sin_theta = sqrt(1 - cos_theta * cos_theta);
 
-
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     if (i >= ghostWidth and i < isize - ghostWidth and j >= ghostWidth and j < jsize - ghostWidth)
     {
 
@@ -927,7 +904,7 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
 
   KOKKOS_INLINE_FUNCTION
   void
-  do_init_energy(const int & index) const
+  do_init_energy(const int & i, const int & j) const
   {
 
     const int isize = params.isize;
@@ -936,9 +913,6 @@ class InitFieldLoopFunctor2D_MHD : public MHDBaseFunctor2D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j;
-    index2coord(index, i, j, isize, jsize);
-
     if (i >= ghostWidth and i < isize - ghostWidth and j >= ghostWidth and j < jsize - ghostWidth)
     {
 
@@ -984,17 +958,20 @@ class InitWaveFunctor2D_MHD : public MHDBaseFunctor2D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, WaveParams wParams, DataArray2d Udata, int nbCells)
+  apply(HydroParams params, WaveParams wParams, DataArray2d Udata)
   {
 
     InitWaveFunctor2D_MHD functor(params, wParams, Udata);
 
-    Kokkos::parallel_for("InitWaveFunctor2D_MHD", Kokkos::RangePolicy<>(0, nbCells), functor);
+    Kokkos::parallel_for(
+      "InitWaveFunctor2D_MHD",
+      Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }),
+      functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j) const
   {} // end operator ()
 
   WaveParams  wParams;
diff --git a/src/muscl/MHDInitFunctors3D.h b/src/muscl/MHDInitFunctors3D.h
index c84ba2c..5e996a4 100644
--- a/src/muscl/MHDInitFunctors3D.h
+++ b/src/muscl/MHDInitFunctors3D.h
@@ -42,20 +42,20 @@ class InitImplodeFunctor3D_MHD : public MHDBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, ImplodeParams iparams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, ImplodeParams iparams, DataArray3d Udata)
   {
     InitImplodeFunctor3D_MHD functor(params, iparams, Udata);
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitImplodeFunctor3D_MHD",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -83,9 +83,6 @@ class InitImplodeFunctor3D_MHD : public MHDBaseFunctor3D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -166,20 +163,20 @@ class InitBlastFunctor3D_MHD : public MHDBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, BlastParams bParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, BlastParams bParams, DataArray3d Udata)
   {
     InitBlastFunctor3D_MHD functor(params, bParams, Udata);
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitBlastFunctor3D_MHD",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -219,9 +216,6 @@ class InitBlastFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t blast_pressure_in = bParams.blast_pressure_in;
     const real_t blast_pressure_out = bParams.blast_pressure_out;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -285,37 +279,40 @@ class InitOrszagTangFunctor3D : public MHDBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, OrszagTangParams otParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, OrszagTangParams otParams, DataArray3d Udata)
   {
     InitOrszagTangFunctor3D functor(params, otParams, Udata);
 
     functor.phase = INIT_ALL_VAR_BUT_ENERGY;
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitOrszagTangFunctor3D",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
 
     functor.phase = INIT_ENERGY;
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitOrszagTangFunctor3D - energy",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
     if (phase == INIT_ALL_VAR_BUT_ENERGY)
-      init_all_var_but_energy(index);
+      init_all_var_but_energy(i, j, k);
     else if (phase == INIT_ENERGY)
-      init_energy(index);
+      init_energy(i, j, k);
 
   } // end operator ()
 
   KOKKOS_INLINE_FUNCTION
   void
-  init_all_var_but_energy(const int index) const
+  init_all_var_but_energy(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -353,9 +350,6 @@ class InitOrszagTangFunctor3D : public MHDBaseFunctor3D
     const double v0 = 1.0;
     const double kt = otParams.kt;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     double xPos = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     double yPos = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     double zPos = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -383,13 +377,13 @@ class InitOrszagTangFunctor3D : public MHDBaseFunctor3D
 
   KOKKOS_INLINE_FUNCTION
   void
-  init_energy(const int index) const
+  init_energy(const int & i, const int & j, const int & k) const
   {
 
     const int isize = params.isize;
     const int jsize = params.jsize;
-    const int ksize = params.ksize;
-    // const int ghostWidth = params.ghostWidth;
+    // const int ksize = params.ksize;
+    //  const int ghostWidth = params.ghostWidth;
 
     const real_t gamma0 = params.settings.gamma0;
 
@@ -398,9 +392,6 @@ class InitOrszagTangFunctor3D : public MHDBaseFunctor3D
     // const double d0    = gamma0*p0;
     // const double v0    = 1.0;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     // double xPos = xmin + dx/2 + (i-ghostWidth)*dx;
     // double yPos = ymin + dy/2 + (j-ghostWidth)*dy;
     // double zPos = zmin + dz/2 + (k-ghostWidth)*dz;
@@ -438,20 +429,20 @@ class InitKelvinHelmholtzFunctor3D_MHD : public MHDBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, KHParams khParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, KHParams khParams, DataArray3d Udata)
   {
     InitKelvinHelmholtzFunctor3D_MHD functor(params, khParams, Udata);
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitKelvinHelmholtzFunctor3D_MHD",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -489,9 +480,6 @@ class InitKelvinHelmholtzFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t ampl = khParams.amplitude;
     const real_t pressure = khParams.pressure;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -647,20 +635,20 @@ class InitRotorFunctor3D_MHD : public MHDBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, RotorParams rParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, RotorParams rParams, DataArray3d Udata)
   {
     InitRotorFunctor3D_MHD functor(params, rParams, Udata);
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitRotorFunctor3D_MHD",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
 #ifdef USE_MPI
@@ -699,9 +687,6 @@ class InitRotorFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t xCenter = (xmax + xmin) / 2;
     const real_t yCenter = (ymax + ymin) / 2;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
 
@@ -790,10 +775,7 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
   struct TagInitEnergy
   {};
 
-  InitFieldLoopFunctor3D_MHD(HydroParams     params,
-                             FieldLoopParams flParams,
-                             DataArray3d     Udata,
-                             int             nbCells)
+  InitFieldLoopFunctor3D_MHD(HydroParams params, FieldLoopParams flParams, DataArray3d Udata)
     : MHDBaseFunctor3D(params)
     , flParams(flParams)
     , Udata(Udata)
@@ -803,47 +785,53 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, FieldLoopParams flParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, FieldLoopParams flParams, DataArray3d Udata)
   {
-    InitFieldLoopFunctor3D_MHD functor(params, flParams, Udata, nbCells);
+    InitFieldLoopFunctor3D_MHD functor(params, flParams, Udata);
 
     functor.phase = COMPUTE_VECTOR_POTENTIAL;
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitFieldLoopFunctor3D_MHD - vector potenetial",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
 
     functor.phase = DO_INIT_CONDITION;
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitFieldLoopFunctor3D_MHD - init",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
 
     functor.phase = DO_INIT_ENERGY;
-    Kokkos::parallel_for(nbCells, functor);
+    Kokkos::parallel_for("InitFieldLoopFunctor3D_MHD - init energy",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         functor);
 
   } // apply
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
     if (phase == COMPUTE_VECTOR_POTENTIAL)
     {
-      compute_vector_potential(index);
+      compute_vector_potential(i, j, k);
     }
     else if (phase == DO_INIT_CONDITION)
     {
-      do_init_condition(index);
+      do_init_condition(i, j, k);
     }
     else if (phase == DO_INIT_ENERGY)
     {
-      do_init_energy(index);
+      do_init_energy(i, j, k);
     }
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  compute_vector_potential(const int & index) const
+  compute_vector_potential(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
     const int nx = params.nx;
@@ -872,9 +860,6 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t radius = flParams.radius;
     const real_t amplitude = flParams.amplitude;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     // real_t z = zmin + dz/2 + (k+nz*k_mpi-ghostWidth)*dz;
@@ -893,7 +878,7 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
 
   KOKKOS_INLINE_FUNCTION
   void
-  do_init_condition(const int & index) const
+  do_init_condition(const int & i, const int & j, const int & k) const
   {
 
     const int isize = params.isize;
@@ -938,11 +923,10 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t cos_theta = 2.0 / sqrt(5.0);
     const real_t sin_theta = sqrt(1 - cos_theta * cos_theta);
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
-    if (i >= ghostWidth and i < isize - ghostWidth and j >= ghostWidth and
-        j < jsize - ghostWidth and k >= ghostWidth and k < ksize - ghostWidth)
+    // clang-format off
+    if (i >= ghostWidth and i < isize - ghostWidth and
+        j >= ghostWidth and j < jsize - ghostWidth and
+        k >= ghostWidth and k < ksize - ghostWidth)
     {
 
       real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
@@ -984,12 +968,13 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
       Udata(i, j, k, IC) = (A(i + 1, j, k, 1) - A(i, j, k, 1)) / dx -
                            (A(i, j + 1, k, 0) - A(i, j, k, 0)) / dy; //+ amp*(drand48()-0.5);
     }
+    // clang-format on
 
   } // end do_init_condition
 
   KOKKOS_INLINE_FUNCTION
   void
-  do_init_energy(const int & index) const
+  do_init_energy(const int & i, const int & j, const int & k) const
   {
 
     const int isize = params.isize;
@@ -999,11 +984,10 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
 
     const real_t gamma0 = params.settings.gamma0;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
-    if (i >= ghostWidth and i < isize - ghostWidth and j >= ghostWidth and
-        j < jsize - ghostWidth and k >= ghostWidth and k < ksize - ghostWidth)
+    // clang-format off
+    if (i >= ghostWidth and i < isize - ghostWidth and
+        j >= ghostWidth and j < jsize - ghostWidth and
+        k >= ghostWidth and k < ksize - ghostWidth)
     {
 
       // total energy
@@ -1024,6 +1008,7 @@ class InitFieldLoopFunctor3D_MHD : public MHDBaseFunctor3D
             Udata(i, j, k, ID);
       }
     }
+    // clang-format on
 
   } // end do_init_energy
 
@@ -1058,7 +1043,7 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
   };
 
 public:
-  InitWaveFunctor3D_MHD(HydroParams params, WaveParams wParams, DataArray3d Udata, int nbCells)
+  InitWaveFunctor3D_MHD(HydroParams params, WaveParams wParams, DataArray3d Udata)
     : MHDBaseFunctor3D(params)
     , wParams(wParams)
     , Udata(Udata)
@@ -1067,48 +1052,54 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
     A = DataArrayVector3("A", params.isize, params.jsize, params.ksize);
 
     phase = COMPUTE_VECTOR_POTENTIAL;
-    Kokkos::parallel_for(nbCells, *this);
+    Kokkos::parallel_for("InitWaveFunctor3D_MHD - vector potential",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         *this);
 
     phase = COMPUTE_FACE_CENTERED_B;
-    Kokkos::parallel_for(nbCells, *this);
+    Kokkos::parallel_for("InitWaveFunctor3D_MHD - face centered mag field init",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         *this);
 
     phase = DO_INIT_CONDITION;
-    Kokkos::parallel_for(nbCells, *this);
+    Kokkos::parallel_for("InitWaveFunctor3D_MHD - init",
+                         Kokkos::MDRangePolicy<Kokkos::Rank<3>>(
+                           { 0, 0, 0 }, { params.isize, params.jsize, params.ksize }),
+                         *this);
   };
 
   // static method which does it all: create and execute functor
   static void
-  apply(HydroParams params, WaveParams wParams, DataArray3d Udata, int nbCells)
+  apply(HydroParams params, WaveParams wParams, DataArray3d Udata)
   {
-    InitWaveFunctor3D_MHD functor(params, wParams, Udata, nbCells);
+    InitWaveFunctor3D_MHD functor(params, wParams, Udata);
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  operator()(const int & index) const
+  operator()(const int & i, const int & j, const int & k) const
   {
     if (phase == COMPUTE_VECTOR_POTENTIAL)
     {
-      compute_vector_potential(index);
+      compute_vector_potential(i, j, k);
     }
     else if (phase == COMPUTE_FACE_CENTERED_B)
     {
-      compute_face_centered_B(index);
+      compute_face_centered_B(i, j, k);
     }
     else if (phase == DO_INIT_CONDITION)
     {
-      do_init_condition(index);
+      do_init_condition(i, j, k);
     }
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  compute_vector_potential(const int & index) const
+  compute_vector_potential(const int & i, const int & j, const int & k) const
   {
 
-    const int isize = params.isize;
-    const int jsize = params.jsize;
-    const int ksize = params.ksize;
     const int ghostWidth = params.ghostWidth;
 
     const int nx = params.nx;
@@ -1133,9 +1124,6 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t dy = params.dy;
     const real_t dz = params.dz;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     real_t x = xmin + dx / 2 + (i + nx * i_mpi - ghostWidth) * dx;
     real_t y = ymin + dy / 2 + (j + ny * j_mpi - ghostWidth) * dy;
     real_t z = zmin + dz / 2 + (k + nz * k_mpi - ghostWidth) * dz;
@@ -1190,7 +1178,7 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
 
   KOKKOS_INLINE_FUNCTION
   void
-  compute_face_centered_B(const int & index) const
+  compute_face_centered_B(const int & i, const int & j, const int & k) const
   {
     const int isize = params.isize;
     const int jsize = params.jsize;
@@ -1201,11 +1189,10 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
     const real_t dy = params.dy;
     const real_t dz = params.dz;
 
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
-    if (i >= ghostWidth - 1 and i < isize - ghostWidth + 1 and j >= ghostWidth - 1 and
-        j < jsize - ghostWidth + 1 and k >= ghostWidth - 1 and k < ksize - ghostWidth + 1)
+    // clang-format off
+    if (i >= ghostWidth - 1 and i < isize - ghostWidth + 1 and
+        j >= ghostWidth - 1 and j < jsize - ghostWidth + 1 and
+        k >= ghostWidth - 1 and k < ksize - ghostWidth + 1)
     {
       // bx
       Udata(i, j, k, IBX) =
@@ -1219,11 +1206,12 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
       Udata(i, j, k, IBZ) =
         (A(i + 1, j, k, 1) - A(i, j, k, 1)) / dx - (A(i, j + 1, k, 0) - A(i, j, k, 0)) / dy;
     }
+    // clang-format on
   }
 
   KOKKOS_INLINE_FUNCTION
   void
-  do_init_condition(const int & index) const
+  do_init_condition(const int & i, const int & j, const int & k) const
   {
 
     const int isize = params.isize;
@@ -1274,10 +1262,6 @@ class InitWaveFunctor3D_MHD : public MHDBaseFunctor3D
 
     const real_t k_par = wParams.k_par;
 
-
-    int i, j, k;
-    index2coord(index, i, j, k, isize, jsize, ksize);
-
     if (i >= ghostWidth and i < isize - ghostWidth and j >= ghostWidth and
         j < jsize - ghostWidth and k >= ghostWidth and k < ksize - ghostWidth)
     {
diff --git a/src/muscl/MHDRunFunctors2D.h b/src/muscl/MHDRunFunctors2D.h
index 5d5a3c4..f44f646 100644
--- a/src/muscl/MHDRunFunctors2D.h
+++ b/src/muscl/MHDRunFunctors2D.h
@@ -196,7 +196,7 @@ class ComputeFluxesAndStoreFunctor2D_MHD : public MHDBaseFunctor2D
   {
     ComputeFluxesAndStoreFunctor2D_MHD functor(
       params, Qm_x, Qm_y, Qp_x, Qp_y, Flux_x, Flux_y, dtdx, dtdy);
-    Kokkos::parallel_for(
+    Kokkos::parallel_for("ComputeFluxesAndStoreFunctor2D_MHD",
       Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }), functor);
   }
 
@@ -289,7 +289,7 @@ class ComputeEmfAndStoreFunctor2D : public MHDBaseFunctor2D
   {
     ComputeEmfAndStoreFunctor2D functor(
       params, QEdge_RT, QEdge_RB, QEdge_LT, QEdge_LB, Emf, dtdx, dtdy);
-    Kokkos::parallel_for(
+    Kokkos::parallel_for("ComputeEmfAndStoreFunctor2D",
       Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }), functor);
   }
 
@@ -394,7 +394,7 @@ class ComputeTraceFunctor2D_MHD : public MHDBaseFunctor2D
                                       QEdge_LB,
                                       dtdx,
                                       dtdy);
-    Kokkos::parallel_for(
+    Kokkos::parallel_for("ComputeTraceFunctor2D_MHD",
       Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }), functor);
   }
 
@@ -492,7 +492,7 @@ class UpdateFunctor2D_MHD : public MHDBaseFunctor2D
         real_t      dtdy)
   {
     UpdateFunctor2D_MHD functor(params, Udata, FluxData_x, FluxData_y, dtdx, dtdy);
-    Kokkos::parallel_for(
+    Kokkos::parallel_for("UpdateFunctor2D_MHD",
       Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }), functor);
   }
 
@@ -590,7 +590,7 @@ class UpdateEmfFunctor2D : public MHDBaseFunctor2D
   apply(HydroParams params, DataArray2d Udata, DataArrayScalar Emf, real_t dtdx, real_t dtdy)
   {
     UpdateEmfFunctor2D functor(params, Udata, Emf, dtdx, dtdy);
-    Kokkos::parallel_for(
+    Kokkos::parallel_for("UpdateEmfFunctor2D",
       Kokkos::MDRangePolicy<Kokkos::Rank<2>>({ 0, 0 }, { params.isize, params.jsize }), functor);
   }
 
diff --git a/src/muscl/SolverHydroMuscl.cpp b/src/muscl/SolverHydroMuscl.cpp
index e0ff4d5..1008845 100644
--- a/src/muscl/SolverHydroMuscl.cpp
+++ b/src/muscl/SolverHydroMuscl.cpp
@@ -87,7 +87,7 @@ SolverHydroMuscl<2>::init_four_quadrant(DataArray Udata)
   primToCons_2D(U2, params.settings.gamma0);
   primToCons_2D(U3, params.settings.gamma0);
 
-  InitFourQuadrantFunctor2D::apply(params, Udata, configNumber, U0, U1, U2, U3, xt, yt, nbCells);
+  InitFourQuadrantFunctor2D::apply(params, Udata, configNumber, U0, U1, U2, U3, xt, yt);
 
 } // SolverHydroMuscl<2>::init_four_quadrant
 
@@ -100,7 +100,7 @@ SolverHydroMuscl<2>::init_isentropic_vortex(DataArray Udata)
 
   IsentropicVortexParams iparams(configMap);
 
-  InitIsentropicVortexFunctor2D::apply(params, iparams, Udata, nbCells);
+  InitIsentropicVortexFunctor2D::apply(params, iparams, Udata);
 
 } // SolverHydroMuscl<2>::init_isentropic_vortex
 
diff --git a/src/muscl/SolverHydroMuscl.h b/src/muscl/SolverHydroMuscl.h
index d1d6c2a..91264e3 100644
--- a/src/muscl/SolverHydroMuscl.h
+++ b/src/muscl/SolverHydroMuscl.h
@@ -162,7 +162,6 @@ class SolverHydroMuscl : public ppkMHD::SolverBase
   save_solution_impl();
 
   int isize, jsize, ksize;
-  int nbCells;
 
 }; // class SolverHydroMuscl
 
@@ -190,15 +189,11 @@ SolverHydroMuscl<dim>::SolverHydroMuscl(HydroParams & params, ConfigMap & config
   , isize(params.isize)
   , jsize(params.jsize)
   , ksize(params.ksize)
-  , nbCells(params.isize * params.jsize)
 {
 
   solver_type = SOLVER_MUSCL_HANCOCK;
 
-  if (dim == 3)
-    nbCells = params.isize * params.jsize * params.ksize;
-
-  m_nCells = nbCells;
+  m_nCells = dim == 2 ? params.isize * params.jsize : params.isize * params.jsize * params.ksize;
   m_nDofsPerCell = 1;
 
   int nbvar = params.nbvar;
@@ -377,7 +372,7 @@ SolverHydroMuscl<dim>::init_implode(DataArray Udata)
     typename std::conditional<dim == 2, InitImplodeFunctor2D, InitImplodeFunctor3D>::type;
 
   // perform init
-  InitImplodeFunctor::apply(params, iparams, Udata, nbCells);
+  InitImplodeFunctor::apply(params, iparams, Udata);
 
 } // SolverHydroMuscl::init_implode
 
@@ -399,7 +394,7 @@ SolverHydroMuscl<dim>::init_blast(DataArray Udata)
     typename std::conditional<dim == 2, InitBlastFunctor2D, InitBlastFunctor3D>::type;
 
   // perform init
-  InitBlastFunctor::apply(params, blastParams, Udata, nbCells);
+  InitBlastFunctor::apply(params, blastParams, Udata);
 
 } // SolverHydroMuscl::init_blast
 
@@ -427,7 +422,7 @@ SolverHydroMuscl<dim>::init_kelvin_helmholtz(DataArray Udata)
     conditional<dim == 2, InitKelvinHelmholtzFunctor2D, InitKelvinHelmholtzFunctor3D>::type;
 
   // perform init
-  InitKelvinHelmholtzFunctor::apply(params, khParams, Udata, nbCells);
+  InitKelvinHelmholtzFunctor::apply(params, khParams, Udata);
 
 } // SolverHydroMuscl::init_kelvin_helmholtz
 
@@ -453,7 +448,7 @@ SolverHydroMuscl<dim>::init_gresho_vortex(DataArray Udata)
     typename std::conditional<dim == 2, InitGreshoVortexFunctor2D, InitGreshoVortexFunctor3D>::type;
 
   // perform init
-  InitGreshoVortexFunctor::apply(params, gvParams, Udata, nbCells);
+  InitGreshoVortexFunctor::apply(params, gvParams, Udata);
 
 } // SolverHydroMuscl<dim>::init_gresho_vortex
 
diff --git a/src/muscl/SolverMHDMuscl.h b/src/muscl/SolverMHDMuscl.h
index b98aa81..516374b 100644
--- a/src/muscl/SolverMHDMuscl.h
+++ b/src/muscl/SolverMHDMuscl.h
@@ -183,7 +183,6 @@ class SolverMHDMuscl : public ppkMHD::SolverBase
   save_solution_impl();
 
   int isize, jsize, ksize;
-  int nbCells;
 
 }; // class SolverMHDMuscl
 
@@ -232,15 +231,11 @@ SolverMHDMuscl<dim>::SolverMHDMuscl(HydroParams & params, ConfigMap & configMap)
   , isize(params.isize)
   , jsize(params.jsize)
   , ksize(params.ksize)
-  , nbCells(params.isize * params.jsize)
 {
 
   solver_type = SOLVER_MUSCL_HANCOCK;
 
-  if (dim == 3)
-    nbCells = params.isize * params.jsize * params.ksize;
-
-  m_nCells = nbCells;
+  m_nCells = dim == 2 ? params.isize * params.jsize : params.isize * params.jsize * params.ksize;
   m_nDofsPerCell = 1;
 
   int nbvar = params.nbvar;
@@ -423,7 +418,7 @@ SolverMHDMuscl<dim>::init_blast(DataArray Udata)
     typename std::conditional<dim == 2, InitBlastFunctor2D_MHD, InitBlastFunctor3D_MHD>::type;
 
   // perform init
-  InitBlastFunctor::apply(params, blastParams, Udata, nbCells);
+  InitBlastFunctor::apply(params, blastParams, Udata);
 
 } // SolverMHDMuscl::init_blast
 
@@ -444,7 +439,7 @@ SolverMHDMuscl<dim>::init_orszag_tang(DataArray Udata)
   using InitOrszagTangFunctor =
     typename std::conditional<dim == 2, InitOrszagTangFunctor2D, InitOrszagTangFunctor3D>::type;
 
-  InitOrszagTangFunctor::apply(params, otParams, Udata, nbCells);
+  InitOrszagTangFunctor::apply(params, otParams, Udata);
 
 } // init_orszag_tang
 
@@ -472,7 +467,7 @@ SolverMHDMuscl<dim>::init_kelvin_helmholtz(DataArray Udata)
     conditional<dim == 2, InitKelvinHelmholtzFunctor2D_MHD, InitKelvinHelmholtzFunctor3D_MHD>::type;
 
   // perform init
-  InitKelvinHelmholtzFunctor::apply(params, khParams, Udata, nbCells);
+  InitKelvinHelmholtzFunctor::apply(params, khParams, Udata);
 
 } // init_kelvin_helmholtz
 
@@ -494,7 +489,7 @@ SolverMHDMuscl<dim>::init_implode(DataArray Udata)
     typename std::conditional<dim == 2, InitImplodeFunctor2D_MHD, InitImplodeFunctor3D_MHD>::type;
 
   // perform init
-  InitImplodeFunctor::apply(params, implodeParams, Udata, nbCells);
+  InitImplodeFunctor::apply(params, implodeParams, Udata);
 
 } // SolverMHDMuscl::init_implode
 
@@ -516,7 +511,7 @@ SolverMHDMuscl<dim>::init_rotor(DataArray Udata)
     typename std::conditional<dim == 2, InitRotorFunctor2D_MHD, InitRotorFunctor3D_MHD>::type;
 
   // perform init
-  InitRotorFunctor::apply(params, rotorParams, Udata, nbCells);
+  InitRotorFunctor::apply(params, rotorParams, Udata);
 
 } // SolverMHDMuscl::init_rotor
 
@@ -538,7 +533,7 @@ SolverMHDMuscl<dim>::init_field_loop(DataArray Udata)
     conditional<dim == 2, InitFieldLoopFunctor2D_MHD, InitFieldLoopFunctor3D_MHD>::type;
 
   // perform init
-  InitFieldLoopFunctor::apply(params, flParams, Udata, nbCells);
+  InitFieldLoopFunctor::apply(params, flParams, Udata);
 
 } // SolverMHDMuscl::init_field_loop
 
@@ -559,7 +554,7 @@ SolverMHDMuscl<dim>::init_wave(DataArray Udata)
     typename std::conditional<dim == 2, InitWaveFunctor2D_MHD, InitWaveFunctor3D_MHD>::type;
 
   // perform init
-  InitWaveFunctor::apply(params, wParams, Udata, nbCells);
+  InitWaveFunctor::apply(params, wParams, Udata);
 
 } // SolverMHDMuscl::init_wave