Twistshift into paralleltransforms

include/bout/paralleltransform.hxx

@@ -79,9 +79,23 @@ public:
   /// Output variables used by a ParallelTransform instance to the dump files
   virtual void outputVars(Datafile& UNUSED(file)) {}
 
-  /// If \p twist_shift_enabled is true, does a `Field3D` with Y direction \p ytype
-  /// require a twist-shift at branch cuts on closed field lines?
-  virtual bool requiresTwistShift(bool twist_shift_enabled, YDirectionType ytype) = 0;
+  /// If \p twist_shift_enabled is true, apply twist-shift to `Field3D f`
+  virtual void applyTwistShift(Field3D& f, bool twist_shift_enabled) = 0;
+
+  /// Shift part of a Field3D by a given angle in Z
+  ///
+  /// @param[inout] var  The variable to be modified in-place
+  /// @param[in] jx      X index
+  /// @param[in] jy      Y index
+  /// @param[in] zangle  The Z angle to apply
+  virtual void shiftZ(Field3D& f, int jx, int jy, BoutReal zangle);
+
+  /// Apply a phase shift by a given angle \p zangle in Z to all points
+  ///
+  /// @param[inout] var  The variable to modify in-place
+  /// @param[in] zangle  The angle to shift by in Z
+  /// @param[in] rgn     The region to calculate the result over
+  void shiftZ(Field3D &var, BoutReal zangle, const std::string& rgn="RGN_ALL");
 
 protected:
   /// This method should be called in the constructor to check that if the grid
@@ -98,7 +112,9 @@ protected:
  */
 class ParallelTransformIdentity : public ParallelTransform {
 public:
-  ParallelTransformIdentity(Mesh& mesh_in) : ParallelTransform(mesh_in) {
+  ParallelTransformIdentity(Mesh& mesh_in, std::vector<BoutReal> ShiftAngle_in)
+      : ParallelTransform(mesh_in), ShiftAngle(std::move(ShiftAngle_in)) {
+
     // check the coordinate system used for the grid data source
     ParallelTransformIdentity::checkInputGrid();
   }
@@ -137,14 +153,12 @@ public:
 
   bool canToFromFieldAligned() override { return true; }
 
-  bool requiresTwistShift(bool twist_shift_enabled, YDirectionType UNUSED(ytype)) override {
-    // All Field3Ds require twist-shift, because all are effectively field-aligned, but
-    // allow twist-shift to be turned off by twist_shift_enabled
-    return twist_shift_enabled;
-  }
+  void applyTwistShift(Field3D& f, bool twist_shift_enabled);
 
 protected:
   void checkInputGrid() override;
+
+  std::vector<BoutReal> ShiftAngle;  ///< Angle for twist-shift location
 };
 
 /*!
@@ -192,14 +206,7 @@ public:
   /// Save zShift to the output
   void outputVars(Datafile& file) override;
 
-  bool requiresTwistShift(bool twist_shift_enabled, YDirectionType ytype) override {
-    // Twist-shift only if field-aligned
-    if (ytype == YDirectionType::Aligned and not twist_shift_enabled) {
-      throw BoutException("'TwistShift = true' is required to communicate field-aligned "
-          "Field3Ds when using ShiftedMetric.");
-    }
-    return ytype == YDirectionType::Aligned;
-  }
+  void applyTwistShift(Field3D& f, bool twist_shift_enabled);
 
 protected:
   void checkInputGrid() override;
@@ -220,7 +227,12 @@ private:
                                    /// orthogonal coordinates to field-aligned coordinates
   Tensor<dcomplex> fromAlignedPhs; ///< Cache of phase shifts for transforming from
                                    /// field-aligned coordinates to X-Z orthogonal
-  /// coordinates
+                                   /// coordinates
+
+  Matrix<dcomplex> twistShiftDownPhs; ///< Cache of phase shifts for twist-shift by
+                                      /// ShiftAngle going down in Y
+  Matrix<dcomplex> twistShiftUpPhs;   ///< Cache of phase shifts for twist-shift by
+                                      /// ShiftAngle going up in Y
 
   /// Helper POD for parallel slice phase shifts
   struct ParallelSlicePhase {

include/field3d.hxx

@@ -297,9 +297,9 @@ class Field3D : public Field, public FieldData {
   Field3D& ynext(int offset);
   const Field3D& ynext(int offset) const;
 
-  /// If \p twist_shift_enabled is true, does this Field3D require a twist-shift at branch
-  /// cuts on closed field lines?
-  bool requiresTwistShift(bool twist_shift_enabled);
+  /// If \p twist_shift_enabled is true, apply twist-shift to this Field3D at branch cuts
+  /// on closed field lines?
+  void applyTwistShift(bool twist_shift_enabled);
 
   /////////////////////////////////////////////////////////
   // Data access
@@ -661,19 +661,19 @@ inline Field3D lowPass(const Field3D &var, int zmax, REGION rgn) {
 /// @param[in] jx      X index
 /// @param[in] jy      Y index
 /// @param[in] zangle  The Z angle to apply
-void shiftZ(Field3D &var, int jx, int jy, double zangle);
+[[gnu::deprecated("Use ParallelTransform::shiftZ instead")]]
+void shiftZ(Field3D& var, int jx, int jy, double zangle);
 
 /// Apply a phase shift by a given angle \p zangle in Z to all points
 ///
 /// @param[inout] var  The variable to modify in-place
 /// @param[in] zangle  The angle to shift by in Z
 /// @param[in] rgn     The region to calculate the result over
-void shiftZ(Field3D &var, BoutReal zangle, const std::string& rgn="RGN_ALL");
-[[gnu::deprecated("Please use shiftZ(const Field3D& var, BoutReal zangle, "
-    "const std::string& region = \"RGN_ALL\") instead")]]
-inline void shiftZ(Field3D &var, BoutReal zangle, REGION rgn) {
-  return shiftZ(var, zangle, toString(rgn));
-}
+[[gnu::deprecated("Use ParallelTransform::shiftZ instead")]]
+void shiftZ(Field3D& var, BoutReal zangle, const std::string& rgn="RGN_ALL");
+[[gnu::deprecated("Please use ParallelTransform::shiftZ(const Field3D& var, BoutReal "
+    "zangle, const std::string& region = \"RGN_ALL\") instead")]]
+void shiftZ(Field3D& var, BoutReal zangle, REGION rgn);
 
 /// Average in the Z direction
 ///

src/field/field3d.cxx

@@ -208,9 +208,8 @@ Field3D &Field3D::ynext(int dir) {
   return const_cast<Field3D&>(static_cast<const Field3D&>(*this).ynext(dir));
 }
 
-bool Field3D::requiresTwistShift(bool twist_shift_enabled) {
-  return getCoordinates()->getParallelTransform().requiresTwistShift(twist_shift_enabled,
-      getDirectionY());
+void Field3D::applyTwistShift(bool twist_shift_enabled) {
+  getCoordinates()->getParallelTransform().applyTwistShift(*this, twist_shift_enabled);
 }
 
 // Not in header because we need to access fieldmesh
@@ -734,49 +733,6 @@ Field3D lowPass(const Field3D &var, int zmax, bool keep_zonal, const std::string
   return result;
 }
 
-/* 
- * Use FFT to shift by an angle in the Z direction
- */
-void shiftZ(Field3D &var, int jx, int jy, double zangle) {
-  TRACE("shiftZ");
-  checkData(var);
-  var.allocate(); // Ensure that var is unique
-  Mesh *localmesh = var.getMesh();
-
-  int ncz = localmesh->LocalNz;
-  if(ncz == 1)
-    return; // Shifting doesn't do anything
-
-  Array<dcomplex> v(ncz/2 + 1);
-
-  rfft(&(var(jx,jy,0)), ncz, v.begin()); // Forward FFT
-
-  BoutReal zlength = var.getCoordinates()->zlength();
-
-  // Apply phase shift
-  for(int jz=1;jz<=ncz/2;jz++) {
-    BoutReal kwave=jz*2.0*PI/zlength; // wave number is 1/[rad]
-    v[jz] *= dcomplex(cos(kwave*zangle) , -sin(kwave*zangle));
-  }
-
-  irfft(v.begin(), ncz, &(var(jx,jy,0))); // Reverse FFT
-}
-
-void shiftZ(Field3D &var, double zangle, const std::string& rgn) {
-  const auto region_str = toString(rgn);
-
-  // Only allow a whitelist of regions for now
-  ASSERT2(region_str == "RGN_ALL" || region_str == "RGN_NOBNDRY" ||
-          region_str == "RGN_NOX" || region_str == "RGN_NOY");
-
-  const Region<Ind2D> &region = var.getRegion2D(region_str);
-
-  // Could be OpenMP if shiftZ(Field3D, int, int, double) didn't throw
-  BOUT_FOR_SERIAL(i, region) {
-    shiftZ(var, i.x(), i.y(), zangle);
-  }
-}
-
 namespace {
   // Internal routine to avoid ugliness with interactions between CHECK
   // levels and UNUSED parameters
@@ -805,6 +761,17 @@ void checkData(const Field3D &f, const std::string& region) {
 }
 #endif
 
+void shiftZ(Field3D& var, int jx, int jy, double zangle) {
+  var.getCoordinates()->getParallelTransform().shiftZ(var, jx, jy, zangle);
+}
+
+void shiftZ(Field3D& var, BoutReal zangle, const std::string& rgn) {
+  var.getCoordinates()->getParallelTransform().shiftZ(var, zangle, rgn);
+}
+void shiftZ(Field3D& var, BoutReal zangle, REGION rgn) {
+  var.getCoordinates()->getParallelTransform().shiftZ(var, zangle, toString(rgn));
+}
+
 Field2D DC(const Field3D &f, const std::string& rgn) {
   TRACE("DC(Field3D)");
 

src/mesh/coordinates.cxx

@@ -1102,7 +1102,14 @@ void Coordinates::setParallelTransform(Options* options) {
 
   if(ptstr == "identity") {
     // Identity method i.e. no transform needed
-    transform = bout::utils::make_unique<ParallelTransformIdentity>(*localmesh);
+
+    std::vector<BoutReal> ShiftAngle(localmesh->LocalNx);
+    for (int x=0; x<localmesh->LocalNx; x++) {
+      localmesh->periodicY(x, ShiftAngle[x]);
+    }
+
+    transform = bout::utils::make_unique<ParallelTransformIdentity>(*localmesh,
+        ShiftAngle);
 
   } else if (ptstr == "shifted") {
     // Shifted metric method

src/mesh/impls/bout/boutmesh.cxx

@@ -1163,41 +1163,7 @@ int BoutMesh::wait(comm_handle handle) {
   // TWIST-SHIFT CONDITION
   // Loop over 3D fields
   for (const auto &var : ch->var_list.field3d()) {
-    if (var->requiresTwistShift(TwistShift)) {
-
-      // Twist-shift only needed for field-aligned fields
-      int jx, jy;
-
-      // Perform Twist-shift using shifting method
-      if (var->getDirectionY() == YDirectionType::Aligned) {
-        // Only variables in field-aligned coordinates need the twist-shift boundary
-        // condition to be applied
-
-        // Lower boundary
-        if (TS_down_in && (DDATA_INDEST != -1)) {
-          for (jx = 0; jx < DDATA_XSPLIT; jx++)
-            for (jy = 0; jy != MYG; jy++)
-              shiftZ(*var, jx, jy, ShiftAngle[jx]);
-        }
-        if (TS_down_out && (DDATA_OUTDEST != -1)) {
-          for (jx = DDATA_XSPLIT; jx < LocalNx; jx++)
-            for (jy = 0; jy != MYG; jy++)
-              shiftZ(*var, jx, jy, ShiftAngle[jx]);
-        }
-
-        // Upper boundary
-        if (TS_up_in && (UDATA_INDEST != -1)) {
-          for (jx = 0; jx < UDATA_XSPLIT; jx++)
-            for (jy = LocalNy - MYG; jy != LocalNy; jy++)
-              shiftZ(*var, jx, jy, -ShiftAngle[jx]);
-        }
-        if (TS_up_out && (UDATA_OUTDEST != -1)) {
-          for (jx = UDATA_XSPLIT; jx < LocalNx; jx++)
-            for (jy = LocalNy - MYG; jy != LocalNy; jy++)
-              shiftZ(*var, jx, jy, -ShiftAngle[jx]);
-        }
-      }
-    }
+    var->applyTwistShift(TwistShift);
   }
 
 #if CHECK > 0
@@ -1884,6 +1850,29 @@ void BoutMesh::topology() {
   if (UDATA_XSPLIT > LocalNx)
     UDATA_XSPLIT = LocalNx;
 
+  // Create Regions for applying twist-shift
+  auto twist_shift_down = Region<Ind2D>();
+  if (TS_down_in && (DDATA_INDEST != -1)) {
+    twist_shift_down += Region<Ind2D>(0, DDATA_XSPLIT - 1, 0, ystart - 1, 0, 0, LocalNy,
+        1, maxregionblocksize);
+  }
+  if (TS_down_out && (DDATA_OUTDEST != -1)) {
+    twist_shift_down += Region<Ind2D>(DDATA_XSPLIT, LocalNx - 1, 0, ystart - 1, 0, 0,
+        LocalNy, 1, maxregionblocksize);
+  }
+  addRegion2D("TwistShiftDown", twist_shift_down);
+
+  auto twist_shift_up = Region<Ind2D>();
+  if (TS_up_in && (UDATA_INDEST != -1)) {
+    twist_shift_up += Region<Ind2D>(0, UDATA_XSPLIT - 1, yend + 1, LocalNy - 1, 0, 0,
+        LocalNy, 1, maxregionblocksize);
+  }
+  if (TS_up_out && (UDATA_OUTDEST != -1)) {
+    twist_shift_up += Region<Ind2D>(UDATA_XSPLIT, LocalNx - 1, yend + 1, LocalNy - 1, 0,
+        0, LocalNy, 1, maxregionblocksize);
+  }
+  addRegion2D("TwistShiftUp", twist_shift_up);
+
   // Print out settings
   output_info.write("\tMYPE_IN_CORE = %d\n", MYPE_IN_CORE);
   output_info.write("\tDXS = %d, DIN = %d. DOUT = %d\n", DDATA_XSPLIT, DDATA_INDEST,

src/mesh/parallel/fci.hxx

@@ -109,11 +109,13 @@ public:
 
   bool canToFromFieldAligned() override { return false; }
 
-  bool requiresTwistShift(bool UNUSED(twist_shift_enabled), MAYBE_UNUSED(YDirectionType ytype)) override {
+  void applyTwistShift(MAYBE_UNUSED(Field3D& f), bool UNUSED(twist_shift_enabled)) {
     // No Field3Ds require twist-shift, because they cannot be field-aligned
-    ASSERT1(ytype == YDirectionType::Standard);
+    ASSERT1(f.getDirectionY() == YDirectionType::Standard);
+  }
 
-    return false;
+  void shiftZ(Field3D& UNUSED(f), int UNUSED(jx), int UNUSED(jy), BoutReal UNUSED(zangle)) {
+    throw BoutException("FCI method cannot shift Field3D by an angle in Z");
   }
 
 protected:

src/mesh/parallel/identity.cxx

@@ -35,3 +35,25 @@ void ParallelTransformIdentity::checkInputGrid() {
   } // else: parallel_transform variable not found in grid input, indicates older input
     //       file so must rely on the user having ensured the type is correct
 }
+
+void ParallelTransformIdentity::applyTwistShift(Field3D& f, bool twist_shift_enabled) {
+  // All Field3Ds require twist-shift, because all are effectively field-aligned, but
+  // allow twist-shift to be turned off by twist_shift_enabled
+  if (twist_shift_enabled) {
+    // Lower boundary
+    // Note "TwistShiftDown" Region is empty if no twist-shift is required at the lower
+    // boundary of this processor
+    BOUT_FOR(i, f.getRegion2D("TwistShiftDown")) {
+      int x = i.x();
+      ParallelTransform::shiftZ(f, x, i.y(), ShiftAngle[x]);
+    }
+
+    // Upper boundary
+    // Note "TwistShiftUp" Region is empty if no twist-shift is required at the upper
+    // boundary of this processor
+    BOUT_FOR(i, f.getRegion2D("TwistShiftUp")) {
+      int x = i.x();
+      ParallelTransform::shiftZ(f, x, i.y(), -ShiftAngle[x]);
+    }
+  }
+}

src/mesh/parallel/makefile

@@ -2,7 +2,7 @@
 BOUT_TOP = ../../..
 
 DIRS            = 
-SOURCEC		= shiftedmetric.cxx fci.cxx identity.cxx
+SOURCEC		= shiftedmetric.cxx fci.cxx identity.cxx paralleltransform.cxx
 TARGET		= lib
 
 include $(BOUT_TOP)/make.config

src/mesh/parallel/paralleltransform.cxx

@@ -0,0 +1,52 @@
+/*
+ * Shared methods for all ParallelTransform classes
+ *
+ */
+
+#include <bout/constants.hxx>
+#include "bout/mesh.hxx"
+#include "bout/paralleltransform.hxx"
+#include <fft.hxx>
+
+/*
+ * Use FFT to shift by an angle in the Z direction
+ */
+void ParallelTransform::shiftZ(Field3D &f, int jx, int jy, double zangle) {
+  TRACE("shiftZ");
+  checkData(f);
+  f.allocate(); // Ensure that f is unique
+  Mesh *localmesh = f.getMesh();
+
+  int ncz = localmesh->LocalNz;
+  if(ncz == 1)
+    return; // Shifting doesn't do anything
+
+  Array<dcomplex> v(ncz/2 + 1);
+
+  rfft(&(f(jx,jy,0)), ncz, v.begin()); // Forward FFT
+
+  BoutReal zlength = f.getCoordinates()->zlength();
+
+  // Apply phase shift
+  for(int jz=1;jz<=ncz/2;jz++) {
+    BoutReal kwave=jz*2.0*PI/zlength; // wave number is 1/[rad]
+    v[jz] *= dcomplex(cos(kwave*zangle) , -sin(kwave*zangle));
+  }
+
+  irfft(v.begin(), ncz, &(f(jx,jy,0))); // Reverse FFT
+}
+
+void ParallelTransform::shiftZ(Field3D &var, double zangle, const std::string& rgn) {
+  const auto region_str = toString(rgn);
+
+  // Only allow a whitelist of regions for now
+  ASSERT2(region_str == "RGN_ALL" || region_str == "RGN_NOBNDRY" ||
+          region_str == "RGN_NOX" || region_str == "RGN_NOY");
+
+  const Region<Ind2D> &region = var.getRegion2D(region_str);
+
+  // Could be OpenMP if shiftZ(Field3D, int, int, double) didn't throw
+  BOUT_FOR_SERIAL(i, region) {
+    shiftZ(var, i.x(), i.y(), zangle);
+  }
+}