Use global Z index functions in ShiftedMetricInterp

Author: ZedThree

src/mesh/parallel/shiftedmetricinterp.cxx

@@ -38,12 +38,19 @@ ShiftedMetricInterp::ShiftedMetricInterp(Mesh& mesh, CELL_LOC location_in,
                                          Options* opt)
     : ParallelTransform(mesh, opt), location(location_in), zShift(std::move(zShift_in)),
       zlength(zlength_in), ydown_index(mesh.ystart) {
+
+  if (mesh.getNZPE() > 1) {
+    throw BoutException("ShiftedMetricInterp only works with 1 processor in Z");
+  }
+
   // check the coordinate system used for the grid data source
   ShiftedMetricInterp::checkInputGrid();
 
   // Allocate space for interpolator cache: y-guard cells in each direction
   parallel_slice_interpolators.resize(mesh.ystart * 2);
 
+  const BoutReal z_factor = static_cast<BoutReal>(mesh.GlobalNzNoBoundaries) / zlength;
+
   // Create the Interpolation objects and set whether they go up or down the
   // magnetic field
   auto& interp_options = options["zinterpolation"];
@@ -62,26 +69,25 @@ ShiftedMetricInterp::ShiftedMetricInterp(Mesh& mesh, CELL_LOC location_in,
 
     // Find the index positions where the magnetic field line intersects the x-z plane
     // y_offset points up
-    Field3D zt_prime_up(&mesh), zt_prime_down(&mesh);
+    Field3D zt_prime_up(&mesh);
+    Field3D zt_prime_down(&mesh);
     zt_prime_up.allocate();
     zt_prime_down.allocate();
 
     for (const auto& i : zt_prime_up.getRegion(RGN_NOY)) {
       // Field line moves in z by an angle zShift(i,j+1)-zShift(i,j) when going
       // from j to j+1, but we want the shift in index-space
-      zt_prime_up[i] = static_cast<BoutReal>(i.z())
-                       + (zShift[i.yp(y_offset + 1)] - zShift[i])
-                             * static_cast<BoutReal>(mesh.GlobalNz) / zlength;
+      zt_prime_up[i] = static_cast<BoutReal>(mesh.getGlobalZIndexNoBoundaries(i.z()))
+                       + ((zShift[i.yp(y_offset + 1)] - zShift[i]) * z_factor);
     }
 
     parallel_slice_interpolators[yup_index + y_offset]->calcWeights(zt_prime_up);
 
     for (const auto& i : zt_prime_down.getRegion(RGN_NOY)) {
       // Field line moves in z by an angle -(zShift(i,j)-zShift(i,j-1)) when going
       // from j to j-1, but we want the shift in index-space
-      zt_prime_down[i] = static_cast<BoutReal>(i.z())
-                         - (zShift[i] - zShift[i.ym(y_offset + 1)])
-                               * static_cast<BoutReal>(mesh.GlobalNz) / zlength;
+      zt_prime_down[i] = static_cast<BoutReal>(mesh.getGlobalZIndexNoBoundaries(i.z()))
+                         - ((zShift[i] - zShift[i.ym(y_offset + 1)]) * z_factor);
     }
 
     parallel_slice_interpolators[ydown_index + y_offset]->calcWeights(zt_prime_down);
@@ -93,15 +99,16 @@ ShiftedMetricInterp::ShiftedMetricInterp(Mesh& mesh, CELL_LOC location_in,
   interp_from_aligned =
       ZInterpolationFactory::getInstance().create(&interp_options, 0, &mesh);
 
-  Field3D zt_prime_to(&mesh), zt_prime_from(&mesh);
+  Field3D zt_prime_to(&mesh);
+  Field3D zt_prime_from(&mesh);
   zt_prime_to.allocate();
   zt_prime_from.allocate();
 
   for (const auto& i : zt_prime_to) {
     // Field line moves in z by an angle zShift(i,j) when going
     // from y0 to y(j), but we want the shift in index-space
-    zt_prime_to[i] = static_cast<BoutReal>(i.z())
-                     + zShift[i] * static_cast<BoutReal>(mesh.GlobalNz) / zlength;
+    zt_prime_to[i] = static_cast<BoutReal>(mesh.getGlobalZIndexNoBoundaries(i.z()))
+                     + (zShift[i] * z_factor);
   }
 
   interp_to_aligned->calcWeights(zt_prime_to);
@@ -110,8 +117,8 @@ ShiftedMetricInterp::ShiftedMetricInterp(Mesh& mesh, CELL_LOC location_in,
     // Field line moves in z by an angle zShift(i,j) when going
     // from y0 to y(j), but we want the shift in index-space.
     // Here we reverse the shift, so subtract zShift
-    zt_prime_from[i] = static_cast<BoutReal>(i.z())
-                       - zShift[i] * static_cast<BoutReal>(mesh.GlobalNz) / zlength;
+    zt_prime_from[i] = static_cast<BoutReal>(mesh.getGlobalZIndexNoBoundaries(i.z()))
+                       - (zShift[i] * z_factor);
   }
 
   interp_from_aligned->calcWeights(zt_prime_from);