Automatically select apropriate parallel slice

include/bout/field.hxx

@@ -133,6 +133,7 @@ public:
   virtual void setRegion(const std::string& UNUSED(region_name)) {}
   virtual void resetRegion() {}
   virtual std::optional<size_t> getRegionID() const { return {}; }
+  virtual int getYoffset() const { return 0; }
 
 private:
   /// Labels for the type of coordinate system this field is defined over
@@ -186,7 +187,7 @@ template <typename T>
 inline T emptyFrom(const T& f) {
   static_assert(bout::utils::is_Field_v<T>, "emptyFrom only works on Fields");
   return T(f.getMesh(), f.getLocation(), {f.getDirectionY(), f.getDirectionZ()},
-           f.getRegionID())
+           f.getRegionID(), f.getYoffset())
       .allocate();
 }
 

include/bout/field2d.hxx

@@ -67,7 +67,7 @@ public:
   Field2D(Mesh* localmesh = nullptr, CELL_LOC location_in = CELL_CENTRE,
           DirectionTypes directions_in = {YDirectionType::Standard,
                                           ZDirectionType::Average},
-          std::optional<size_t> region = {});
+          std::optional<size_t> region = {}, int yoffset = 0);
 
   /*!
    * Copy constructor. After this both fields

include/bout/field3d.hxx

@@ -29,6 +29,7 @@ class Field3D;
 #include "bout/array.hxx"
 #include "bout/assert.hxx"
 #include "bout/bout_types.hxx"
+#include "bout/boutexception.hxx"
 #include "bout/field.hxx"
 #include "bout/field2d.hxx"
 #include "bout/fieldperp.hxx"
@@ -167,7 +168,7 @@ public:
   Field3D(Mesh* localmesh = nullptr, CELL_LOC location_in = CELL_CENTRE,
           DirectionTypes directions_in = {YDirectionType::Standard,
                                           ZDirectionType::Standard},
-          std::optional<size_t> regionID = {});
+          std::optional<size_t> regionID = {}, int yoffset = 0);
 
   /*!
    * Copy constructor
@@ -271,23 +272,27 @@ public:
   /// Return reference to yup field
   Field3D& yup(std::vector<Field3D>::size_type index = 0) {
     ASSERT2(index < yup_fields.size());
+    ASSERT2(yup_fields[index].yoffset == static_cast<int>(index) + 1);
     return yup_fields[index];
   }
   /// Return const reference to yup field
   const Field3D& yup(std::vector<Field3D>::size_type index = 0) const {
     ASSERT2(index < yup_fields.size());
+    ASSERT2(yup_fields[index].yoffset == static_cast<int>(index) + 1);
     return yup_fields[index];
   }
 
   /// Return reference to ydown field
   Field3D& ydown(std::vector<Field3D>::size_type index = 0) {
     ASSERT2(index < ydown_fields.size());
+    ASSERT2(ydown_fields[index].yoffset == -1 - static_cast<int>(index));
     return ydown_fields[index];
   }
 
   /// Return const reference to ydown field
   const Field3D& ydown(std::vector<Field3D>::size_type index = 0) const {
     ASSERT2(index < ydown_fields.size());
+    ASSERT2(ydown_fields[index].yoffset == -1 - static_cast<int>(index));
     return ydown_fields[index];
   }
 
@@ -339,6 +344,8 @@ public:
   void setRegion(std::optional<size_t> id) override;
   std::optional<size_t> getRegionID() const override { return regionID; };
 
+  int getYoffset() const override { return yoffset; };
+
   /// Return a Region<Ind2D> reference to use to iterate over the x- and
   /// y-indices of this field
   const Region<Ind2D>& getRegion2D(REGION region) const;
@@ -351,8 +358,27 @@ public:
     return std::end(getRegion("RGN_ALL"));
   };
 
-  BoutReal& operator[](const Ind3D& d) { return data[d.ind]; }
-  const BoutReal& operator[](const Ind3D& d) const { return data[d.ind]; }
+  BoutReal& operator[](const Ind3D& d) {
+    if (d.yoffset != 0) {
+      if (yoffset == 0) {
+        if (hasParallelSlices()) {
+          return ynext(d.yoffset)[d];
+        }
+#if CHECK >= 2
+        if (isFci()) { // We probably should assert here that this is field aligned
+          throw BoutException(
+              "Tried to access parallel slices, but they are not calculated!");
+        }
+#endif
+      } else {
+        ASSERT2(d.yoffset == yoffset);
+      }
+    }
+    return data[d.ind];
+  }
+  const BoutReal& operator[](const Ind3D& d) const {
+    return (*const_cast<Field3D*>(this))[d];
+  }
 
   BoutReal& operator()(const IndPerp& d, int jy);
   const BoutReal& operator()(const IndPerp& d, int jy) const;
@@ -547,6 +573,10 @@ private:
   template <class T>
   Options* track(const T& change, std::string operation);
   Options* track(const BoutReal& change, std::string operation);
+
+public:
+  int yoffset{0};
+  void setParallelRegions();
 };
 
 // Non-member overloaded operators

include/bout/fieldperp.hxx

@@ -59,7 +59,7 @@ public:
             int yindex_in = -1,
             DirectionTypes directions_in = {YDirectionType::Standard,
                                             ZDirectionType::Standard},
-            std::optional<size_t> regionID = {});
+            std::optional<size_t> regionID = {}, int yoffset = 0);
 
   /*!
    * Copy constructor. After this the data

include/bout/region.hxx

@@ -168,9 +168,12 @@ template <IND_TYPE N>
 struct SpecificInd {
   int ind = -1;         ///< 1D index into Field
   int ny = -1, nz = -1; ///< Sizes of y and z dimensions
+  int yoffset = 0;
 
   SpecificInd() = default;
   SpecificInd(int i, int ny, int nz) : ind(i), ny(ny), nz(nz){};
+  SpecificInd(int i, int ny, int nz, int yoffset)
+      : ind(i), ny(ny), nz(nz), yoffset(yoffset){};
   explicit SpecificInd(int i) : ind(i){};
 
   /// Allow explicit conversion to an int
@@ -275,7 +278,9 @@ struct SpecificInd {
     }
   }
 
-  inline SpecificInd xp(int dx = 1) const { return {ind + (dx * ny * nz), ny, nz}; }
+  inline SpecificInd xp(int dx = 1) const {
+    return {ind + (dx * ny * nz), ny, nz, yoffset};
+  }
   /// The index one point -1 in x
   inline SpecificInd xm(int dx = 1) const { return xp(-dx); }
   /// The index one point +1 in y
@@ -286,8 +291,18 @@ struct SpecificInd {
     }
 #endif
     ASSERT3(std::abs(dy) < ny);
-    return {ind + (dy * nz), ny, nz};
+    return {ind + (dy * nz), ny, nz, yoffset + dy};
   }
+  inline SpecificInd yp_no_parallel_shift(int dy = 1) const {
+#if CHECK >= 4
+    if (y() + dy < 0 or y() + dy >= ny) {
+      throw BoutException("Offset in y ({:d}) would go out of bounds at {:d}", dy, ind);
+    }
+#endif
+    ASSERT3(std::abs(dy) < ny);
+    return {ind + (dy * nz), ny, nz, yoffset};
+  }
+
   /// The index one point -1 in y
   inline SpecificInd ym(int dy = 1) const { return yp(-dy); }
   /// The index one point +1 in z. Wraps around zend to zstart
@@ -297,7 +312,7 @@ struct SpecificInd {
   inline SpecificInd zp(int dz = 1) const {
     ASSERT3(dz >= 0);
     dz = dz <= nz ? dz : dz % nz; //Fix in case dz > nz, if not force it to be in range
-    return {(ind + dz) % nz < dz ? ind - nz + dz : ind + dz, ny, nz};
+    return {(ind + dz) % nz < dz ? ind - nz + dz : ind + dz, ny, nz, yoffset};
   }
   /// The index one point -1 in z. Wraps around zstart to zend
   /// An alternative, non-branching calculation is :
@@ -306,7 +321,7 @@ struct SpecificInd {
   inline SpecificInd zm(int dz = 1) const {
     dz = dz <= nz ? dz : dz % nz; //Fix in case dz > nz, if not force it to be in range
     ASSERT3(dz >= 0);
-    return {(ind) % nz < dz ? ind + nz - dz : ind - dz, ny, nz};
+    return {(ind) % nz < dz ? ind + nz - dz : ind - dz, ny, nz, yoffset};
   }
   /// Automatically select zm or zp depending on sign
   inline SpecificInd zpm(int dz) const { return dz > 0 ? zp(dz) : zm(-dz); }

src/field/field2d.cxx

@@ -49,7 +49,7 @@
 #include <bout/assert.hxx>
 
 Field2D::Field2D(Mesh* localmesh, CELL_LOC location_in, DirectionTypes directions_in,
-                 std::optional<size_t> UNUSED(regionID))
+                 std::optional<size_t> UNUSED(regionID), int UNUSED(yoffset))
     : Field(localmesh, location_in, directions_in) {
 
   if (fieldmesh) {

src/field/field3d.cxx

@@ -49,8 +49,8 @@
 
 /// Constructor
 Field3D::Field3D(Mesh* localmesh, CELL_LOC location_in, DirectionTypes directions_in,
-                 std::optional<size_t> regionID)
-    : Field(localmesh, location_in, directions_in), regionID{regionID} {
+                 std::optional<size_t> regionID, const int yoffset)
+    : Field(localmesh, location_in, directions_in), regionID{regionID}, yoffset{yoffset} {
 #if BOUT_USE_TRACK
   name = "<F3D>";
 #endif
@@ -66,7 +66,7 @@ Field3D::Field3D(Mesh* localmesh, CELL_LOC location_in, DirectionTypes direction
 /// later)
 Field3D::Field3D(const Field3D& f)
     : Field(f), data(f.data), yup_fields(f.yup_fields), ydown_fields(f.ydown_fields),
-      regionID(f.regionID) {
+      regionID(f.regionID), yoffset(f.yoffset) {
 
   TRACE("Field3D(Field3D&)");
 
@@ -157,7 +157,9 @@ void Field3D::splitParallelSlices() {
     // Note the fields constructed here will be fully overwritten by the
     // ParallelTransform, so we don't need a full constructor
     yup_fields.emplace_back(fieldmesh);
+    yup_fields[i].yoffset = i + 1;
     ydown_fields.emplace_back(fieldmesh);
+    ydown_fields[i].yoffset = -1 - i;
     if (isFci()) {
       yup_fields[i].setRegion(fmt::format("RGN_YPAR_{:+d}", i + 1));
       ydown_fields[i].setRegion(fmt::format("RGN_YPAR_{:+d}", -i - 1));
@@ -224,6 +226,15 @@ bool Field3D::requiresTwistShift(bool twist_shift_enabled) {
                                                                      getDirectionY());
 }
 
+void Field3D::setParallelRegions() {
+  for (int i = 0; i < fieldmesh->ystart; ++i) {
+    yup_fields[i].yoffset = i + 1;
+    yup_fields[i].setRegion(fmt::format("RGN_YPAR_{:+d}", i + 1));
+    ydown_fields[i].yoffset = -1 - i;
+    ydown_fields[i].setRegion(fmt::format("RGN_YPAR_{:+d}", -1 - i));
+  }
+}
+
 // Not in header because we need to access fieldmesh
 BoutReal& Field3D::operator()(const IndPerp& d, int jy) {
   return operator[](fieldmesh->indPerpto3D(d, jy));
@@ -277,6 +288,7 @@ Field3D& Field3D::operator=(const Field3D& rhs) {
   yup_fields = rhs.yup_fields;
   ydown_fields = rhs.ydown_fields;
   regionID = rhs.regionID;
+  yoffset = rhs.yoffset;
 
   // Copy the data and data sizes
   nx = rhs.nx;
@@ -295,12 +307,13 @@ Field3D& Field3D::operator=(Field3D&& rhs) {
   // Move parallel slices or delete existing ones.
   yup_fields = std::move(rhs.yup_fields);
   ydown_fields = std::move(rhs.ydown_fields);
+  regionID = rhs.regionID;
+  yoffset = rhs.yoffset;
 
   // Move the data and data sizes
   nx = rhs.nx;
   ny = rhs.ny;
   nz = rhs.nz;
-  regionID = rhs.regionID;
 
   data = std::move(rhs.data);
 
@@ -896,6 +909,7 @@ void swap(Field3D& first, Field3D& second) noexcept {
   swap(first.deriv, second.deriv);
   swap(first.yup_fields, second.yup_fields);
   swap(first.ydown_fields, second.ydown_fields);
+  swap(first.yoffset, second.yoffset);
 }
 
 const Region<Ind3D>&

src/field/fieldperp.cxx

@@ -35,7 +35,8 @@
 #include <bout/utils.hxx>
 
 FieldPerp::FieldPerp(Mesh* localmesh, CELL_LOC location_in, int yindex_in,
-                     DirectionTypes directions, std::optional<size_t> UNUSED(regionID))
+                     DirectionTypes directions, std::optional<size_t> UNUSED(regionID),
+                     int UNUSED(yoffset))
     : Field(localmesh, location_in, directions), yindex(yindex_in) {
   if (fieldmesh) {
     nx = fieldmesh->LocalNx;

src/field/gen_fieldops.jinja

@@ -12,6 +12,8 @@
     {% if lhs == rhs == "Field3D" %}
       {{out.name}}.setRegion({{lhs.name}}.getMesh()->getCommonRegion({{lhs.name}}.getRegionID(),
                                                                      {{rhs.name}}.getRegionID()));
+      ASSERT2({{lhs.name}}.yoffset == {{rhs.name}}.yoffset);
+      {{out.name}}.yoffset = {{lhs.name}}.yoffset;
 #if BOUT_USE_FCI_AUTOMAGIC
       if ({{lhs.name}}.isFci() and {{lhs.name}}.hasParallelSlices() and {{rhs.name}}.hasParallelSlices()) {
         {{out.name}}.splitParallelSlices();
@@ -23,6 +25,7 @@
 #endif
     {% elif lhs == "Field3D" %}
       {{out.name}}.setRegion({{lhs.name}}.getRegionID());
+      {{out.name}}.yoffset = {{lhs.name}}.yoffset;
       {% if rhs == "BoutReal" %}
 #if BOUT_USE_FCI_AUTOMAGIC
       if ({{lhs.name}}.isFci() and {{lhs.name}}.hasParallelSlices()) {
@@ -36,6 +39,7 @@
       {% endif %}
     {% elif rhs == "Field3D" %}
       {{out.name}}.setRegion({{rhs.name}}.getRegionID());
+      {{out.name}}.yoffset = {{rhs.name}}.yoffset;
       {% if lhs == "BoutReal" %}
 #if BOUT_USE_FCI_AUTOMAGIC
       if ({{rhs.name}}.isFci() and {{rhs.name}}.hasParallelSlices()) {

src/field/gen_fieldops.py

@@ -65,6 +65,7 @@ def smart_open(filename, mode="r"):
 )
 
 header = """// This file is autogenerated - see gen_fieldops.py
+#include <bout/assert.hxx>
 #include <bout/field2d.hxx>
 #include <bout/field3d.hxx>
 #include <bout/globals.hxx>