WIP: Add communication in Z

The actual significant changes are quite small, but these have big knock-on
effects in the tests. Main changes:

- Z is still periodic, but this is now implemented using guard cells instead of
  modulus in `SpecificInd::zp/zm`
- Adds a new `Mesh` method `sendZ`
- `NZPE` is now a user-settable option
- Domain decomposition is done in Z before the usual X-Y decomposition
- The X/Y _boundary_ regions do not include the Z guards
  - These corners should be filled in by `Mesh::communicate`
- Handle Z guards in `Gridfile`
- Handle Z guards in FFTs for `LaplaceCyclic`, `LaplacePCR`, `LaplacePCR_THOMAS`
  - Pull out common classes for the FFT and DST transforms

Other bug fixes:
- `BoutMesh::GlobalZ` should start at the first interior point, not in the guards

Changes to the tests:
- `FakeMesh` now has 1 Z guard cell, which changes the size of the interior
  domain, and requires many of the `Field*` tests to be updated
- `IndexOffsetTest` needs to skip the Z guards, just like X, Y
- Derivative tests need extra point for interior + 2 * guards
- The Petsc and Hypre Laplace unit tests require the Z guards to be communicated
  in order to fill in the global indices. This requires `FakeMesh` to actually
  implement `sendZ`. This just copies the relevant points into the Z guards
  - Also, we should not check the result in the Z guards
  - I've also pulled out a bunch of common machinery between these two tests
- `BoutMeshTest` now needs to understand Z guards and Z processors

This change also requires `boutdata.collect` to properly handle Z guards as well

Author: ZedThree

CMakeLists.txt

@@ -208,6 +208,8 @@ set(BOUT_SOURCES
   ./src/field/where.cxx
   ./src/invert/fft_fftw.cxx
   ./src/invert/lapack_routines.cxx
+  ./src/invert/laplace/common_transform.cxx
+  ./src/invert/laplace/common_transform.hxx
   ./src/invert/laplace/impls/cyclic/cyclic_laplace.cxx
   ./src/invert/laplace/impls/cyclic/cyclic_laplace.hxx
   ./src/invert/laplace/impls/iterative_parallel_tri/iterative_parallel_tri.cxx

externalpackages/boutdata

@@ -124,7 +124,7 @@ Array<BoutReal> irfft(const Array<dcomplex>& in, int length);
 void assertZSerial(const Mesh& mesh, std::string_view name);
 #else
 inline void assertZSerial([[maybe_unused]] const Mesh& mesh,
-                         [[maybe_unused]] std::string_view name) {}
+                          [[maybe_unused]] std::string_view name) {}
 #endif
 } // namespace fft
 } // namespace bout

include/bout/fft.hxx

@@ -129,20 +129,38 @@ public:
 private:
   Options data;
   std::string filename;
+
+  // These next three quantities are already read in `BoutMesh`, but this way we
+  // don't need the Mesh API to get them here
+
+  // Number of y-boundary guard cells saved in the grid file
   int grid_yguards{0};
   int ny_inner{0};
+  // Number of z-boundary guard cells save in the grid file.
+  // Always zero for older grid files
+  int grid_zguards{0};
 
   bool readgrid_3dvar_fft(Mesh* m, const std::string& name, int yread, int ydest,
                           int ysize, int xread, int xdest, int xsize, Field3D& var);
 
+  // \param yread  Start reading at global y-index
+  // \param ydest  Insert data starting from y=yd
+  // \param ysize  Length of data in Y
+  // \param xread  Start reading at global x-index
+  // \param xdest  Insert data starting from x=xd
+  // \param xsize  Length of data in X
+  // \param zread  Start reading at global z-index
+  // \param zdest  Insert data starting from z=zd
+  // \param zsize  Length of data in Z
   bool readgrid_3dvar_real(const std::string& name, int yread, int ydest, int ysize,
-                           int xread, int xdest, int xsize, Field3D& var);
+                           int xread, int xdest, int xsize, int zread, int zdest,
+                           int zsize, Field3D& var);
 
   bool readgrid_perpvar_fft(Mesh* m, const std::string& name, int xread, int xdest,
                             int xsize, FieldPerp& var);
 
   bool readgrid_perpvar_real(const std::string& name, int xread, int xdest, int xsize,
-                             FieldPerp& var);
+                             int zread, int zdest, int zsize, FieldPerp& var);
 
   // convenience template method to remove code duplication between Field2D,
   // Field3D and FieldPerp versions of get
@@ -151,13 +169,16 @@ private:
                 CELL_LOC location);
   // utility method for Field2D to implement unshared parts of getField
   void readField(Mesh* m, const std::string& name, int ys, int yd, int ny_to_read, int xs,
-                 int xd, int nx_to_read, const std::vector<int>& size, Field2D& var);
+                 int xd, int nx_to_read, int zs, int zd, int nz_to_read,
+                 const std::vector<int>& size, Field2D& var);
   // utility method for Field3D to implement unshared parts of getField
   void readField(Mesh* m, const std::string& name, int ys, int yd, int ny_to_read, int xs,
-                 int xd, int nx_to_read, const std::vector<int>& size, Field3D& var);
+                 int xd, int nx_to_read, int zs, int zd, int nz_to_read,
+                 const std::vector<int>& size, Field3D& var);
   // utility method for FieldPerp to implement unshared parts of getField
   void readField(Mesh* m, const std::string& name, int ys, int yd, int ny_to_read, int xs,
-                 int xd, int nx_to_read, const std::vector<int>& size, FieldPerp& var);
+                 int xd, int nx_to_read, int zs, int zd, int nz_to_read,
+                 const std::vector<int>& size, FieldPerp& var);
 };
 
 /*!

include/bout/griddata.hxx

@@ -50,6 +50,8 @@ class Laplacian;
 
 #include "bout/dcomplex.hxx"
 
+class DSTTransform;
+class FFTTransform;
 class Solver;
 
 constexpr auto LAPLACE_SPT = "spt";
@@ -258,7 +260,7 @@ public:
   /// Coefficients in tridiagonal inversion
   void tridagCoefs(int jx, int jy, int jz, dcomplex& a, dcomplex& b, dcomplex& c,
                    const Field2D* ccoef = nullptr, const Field2D* d = nullptr,
-                   CELL_LOC loc = CELL_DEFAULT);
+                   CELL_LOC loc = CELL_DEFAULT) const;
 
   /*!
    * Create a new Laplacian solver
@@ -301,6 +303,10 @@ public:
   void savePerformance(Solver& solver, const std::string& name);
 
 protected:
+  // Give access for tridagMatrix
+  friend class DSTTransform;
+  friend class FFTTransform;
+
   bool async_send; ///< If true, use asyncronous send in parallel algorithms
 
   int maxmode; ///< The maximum Z mode to solve for
@@ -332,23 +338,24 @@ protected:
 
   void tridagCoefs(int jx, int jy, BoutReal kwave, dcomplex& a, dcomplex& b, dcomplex& c,
                    const Field2D* ccoef = nullptr, const Field2D* d = nullptr,
-                   CELL_LOC loc = CELL_DEFAULT) {
+                   CELL_LOC loc = CELL_DEFAULT) const {
     tridagCoefs(jx, jy, kwave, a, b, c, ccoef, ccoef, d, loc);
   }
   void tridagCoefs(int jx, int jy, BoutReal kwave, dcomplex& a, dcomplex& b, dcomplex& c,
                    const Field2D* c1coef, const Field2D* c2coef, const Field2D* d,
-                   CELL_LOC loc = CELL_DEFAULT);
+                   CELL_LOC loc = CELL_DEFAULT) const;
 
   void tridagMatrix(dcomplex* avec, dcomplex* bvec, dcomplex* cvec, dcomplex* bk, int jy,
                     int kz, BoutReal kwave, const Field2D* a, const Field2D* ccoef,
-                    const Field2D* d, bool includeguards = true, bool zperiodic = true) {
+                    const Field2D* d, bool includeguards = true,
+                    bool zperiodic = true) const {
     tridagMatrix(avec, bvec, cvec, bk, jy, kz, kwave, a, ccoef, ccoef, d, includeguards,
                  zperiodic);
   }
   void tridagMatrix(dcomplex* avec, dcomplex* bvec, dcomplex* cvec, dcomplex* bk, int jy,
                     int kz, BoutReal kwave, const Field2D* a, const Field2D* c1coef,
                     const Field2D* c2coef, const Field2D* d, bool includeguards = true,
-                    bool zperiodic = true);
+                    bool zperiodic = true) const;
   CELL_LOC location;   ///< staggered grid location of this solver
   Mesh* localmesh;     ///< Mesh object for this solver
   Coordinates* coords; ///< Coordinates object, so we only have to call

include/bout/invert_laplace.hxx

@@ -349,6 +349,9 @@ public:
   /// Send only the y-guard cells
   virtual comm_handle sendY(FieldGroup& g, comm_handle handle = nullptr) = 0;
 
+  /// Send only the z-guard cells
+  virtual comm_handle sendZ(FieldGroup& g, comm_handle handle = nullptr) = 0;
+
   /// Wait for the handle, return error code
   virtual int wait(comm_handle handle) = 0; ///< Wait for the handle, return error code
 
@@ -812,9 +815,9 @@ protected:
   /// Read a 1D array of integers
   const std::vector<int> readInts(const std::string& name, int n);
 
-  /// Calculates the size of a message for a given x and y range
-  int msg_len(const std::vector<FieldData*>& var_list, int xge, int xlt, int yge,
-              int ylt);
+  /// Calculates the size of a message for a given (x, y, z) range
+  int msg_len(const std::vector<FieldData*>& var_list, int xge, int xlt, int yge, int ylt,
+              int zge, int zlt);
 
   /// Initialise derivatives
   void derivs_init(Options* options);

include/bout/mesh.hxx

@@ -170,8 +170,8 @@ struct SpecificInd {
   int ny = -1, nz = -1; ///< Sizes of y and z dimensions
 
   SpecificInd() = default;
-  SpecificInd(int i, int ny, int nz) : ind(i), ny(ny), nz(nz){};
-  explicit SpecificInd(int i) : ind(i){};
+  SpecificInd(int i, int ny, int nz) : ind(i), ny(ny), nz(nz) {};
+  explicit SpecificInd(int i) : ind(i) {};
 
   /// Allow explicit conversion to an int
   explicit operator int() const { return ind; }
@@ -290,26 +290,15 @@ struct SpecificInd {
   }
   /// 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
-  /// An alternative, non-branching calculation is :
-  /// ind + dz - nz * ((ind + dz) / nz  - ind / nz)
-  /// but this appears no faster (and perhaps slower).
+  /// The index one point +1 in z
   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};
-  }
-  /// The index one point -1 in z. Wraps around zstart to zend
-  /// An alternative, non-branching calculation is :
-  /// ind - dz + nz * ( (nz + ind) / nz - (nz + ind - dz) / nz)
-  /// but this appears no faster (and perhaps slower).
-  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};
+    if constexpr (N == IND_TYPE::IND_2D) {
+      return *this;
+    }
+    return {ind + dz, ny, nz};
   }
-  /// Automatically select zm or zp depending on sign
-  inline SpecificInd zpm(int dz) const { return dz > 0 ? zp(dz) : zm(-dz); }
+  /// The index one point -1 in z
+  inline SpecificInd zm(int dz = 1) const { return zp(-dz); }
 
   // and for 2 cells
   inline SpecificInd xpp() const { return xp(2); }
@@ -320,9 +309,7 @@ struct SpecificInd {
   inline SpecificInd zmm() const { return zm(2); }
 
   /// Generic offset of \p index in multiple directions simultaneously
-  inline SpecificInd offset(int dx, int dy, int dz) const {
-    return zpm(dz).yp(dy).xp(dx);
-  }
+  inline SpecificInd offset(int dx, int dy, int dz) const { return zp(dz).yp(dy).xp(dx); }
 };
 
 /// Relational operators
@@ -490,10 +477,9 @@ template <typename T = Ind3D>
 class Region {
   // Following prevents a Region being created with anything other
   // than Ind2D, Ind3D or IndPerp as template type
-  static_assert(
-      std::is_base_of_v<
-          Ind2D, T> || std::is_base_of_v<Ind3D, T> || std::is_base_of_v<IndPerp, T>,
-      "Region must be templated with one of IndPerp, Ind2D or Ind3D");
+  static_assert(std::is_base_of_v<Ind2D, T> || std::is_base_of_v<Ind3D, T>
+                    || std::is_base_of_v<IndPerp, T>,
+                "Region must be templated with one of IndPerp, Ind2D or Ind3D");
 
 public:
   using data_type = T;
@@ -569,7 +555,7 @@ public:
   };
 
   Region(RegionIndices& indices, int maxregionblocksize = MAXREGIONBLOCKSIZE)
-      : indices(indices), blocks(getContiguousBlocks(maxregionblocksize)){};
+      : indices(indices), blocks(getContiguousBlocks(maxregionblocksize)) {};
 
   // We need to first set the blocks, and only after that call getRegionIndices.
   // Do not put in the member initialisation