ShiftedMetric with 2 yup/ydown points

include/bout/mesh.hxx

@@ -614,12 +614,12 @@ class Mesh {
   typedef BoutReal (*flux_func)(stencil&, stencil &);
 
   /// Transform a field into field-aligned coordinates
-  const Field3D toFieldAligned(const Field3D &f) {
-    return getParallelTransform().toFieldAligned(f);
+  const Field3D toFieldAligned(const Field3D &f, const REGION region = RGN_NOX) {
+    return getParallelTransform().toFieldAligned(f, region);
   }
   /// Convert back into standard form
-  const Field3D fromFieldAligned(const Field3D &f) {
-    return getParallelTransform().fromFieldAligned(f);
+  const Field3D fromFieldAligned(const Field3D &f, const REGION region = RGN_NOX) {
+    return getParallelTransform().fromFieldAligned(f, region);
   }
 
   bool canToFromFieldAligned() {

include/bout/paralleltransform.hxx

@@ -6,10 +6,13 @@
 #ifndef __PARALLELTRANSFORM_H__
 #define __PARALLELTRANSFORM_H__
 
+#include <datafile.hxx>
 #include <field3d.hxx>
+#include <flexible.hxx>
 #include <boutexception.hxx>
 #include <dcomplex.hxx>
 #include <unused.hxx>
+#include <utils.hxx>
 
 class Mesh;
 
@@ -38,13 +41,16 @@ public:
 
   /// Convert a 3D field into field-aligned coordinates
   /// so that the y index is along the magnetic field
-  virtual const Field3D toFieldAligned(const Field3D &f) = 0;
+  virtual const Field3D toFieldAligned(const Field3D &f, const REGION region = RGN_NOX) = 0;
 
   /// Convert back from field-aligned coordinates
   /// into standard form
-  virtual const Field3D fromFieldAligned(const Field3D &f) = 0;
+  virtual const Field3D fromFieldAligned(const Field3D &f, const REGION region = RGN_NOX) = 0;
 
   virtual bool canToFromFieldAligned() = 0;
+
+  /// Write out ParallelTransform variables to file
+  virtual void outputVars(Datafile &UNUSED(file)) {};
 };
 
 
@@ -65,21 +71,24 @@ public:
    * The field is already aligned in Y, so this
    * does nothing
    */ 
-  const Field3D toFieldAligned(const Field3D &f) override {
+  const Field3D toFieldAligned(const Field3D &f, const REGION UNUSED(region)) override {
     return f;
   }
 
   /*!
    * The field is already aligned in Y, so this
    * does nothing
    */
-  const Field3D fromFieldAligned(const Field3D &f) override {
+  const Field3D fromFieldAligned(const Field3D &f, const REGION UNUSED(region)) override {
     return f;
   }
 
   bool canToFromFieldAligned() override{
     return true;
   }
+
+  /// Write out ParallelTransform variables to file
+  virtual void outputVars(Datafile &UNUSED(file)) {};
 };
 
 /*!
@@ -108,42 +117,68 @@ public:
    * in X-Z, and the metric tensor will need to be changed 
    * if X derivatives are used.
    */
-  const Field3D toFieldAligned(const Field3D &f) override;
+  const Field3D toFieldAligned(const Field3D &f, const REGION region=RGN_NOX) override;
 
   /*!
    * Converts a field back to X-Z orthogonal coordinates
    * from field aligned coordinates.
    */
-  const Field3D fromFieldAligned(const Field3D &f) override;
+  const Field3D fromFieldAligned(const Field3D &f, const REGION region=RGN_NOX) override;
 
   bool canToFromFieldAligned() override{
     return true;
   }
 
-  /// A 3D array, implemented as nested vectors
-  typedef std::vector<std::vector<std::vector<dcomplex>>> arr3Dvec;
 private:
   ShiftedMetric();
 
   Mesh &mesh; ///< The mesh this paralleltransform is part of
 
   /// This is the shift in toroidal angle (z) which takes a point from
   /// X-Z orthogonal to field-aligned along Y.
-  Field2D zShift;
-  std::vector<dcomplex> cmplx; ///< A temporary array, used for input/output to fft routines
-  std::vector<dcomplex> cmplxLoc; ///< A temporary array, used for input/output to fft routines
-
-  arr3Dvec toAlignedPhs; ///< Cache of phase shifts for transforming from X-Z orthogonal coordinates to field-aligned coordinates
-  arr3Dvec fromAlignedPhs; ///< Cache of phase shifts for transforming from field-aligned coordinates to X-Z orthogonal coordinates
-
-  arr3Dvec yupPhs; ///< Cache of phase shifts for calculating yup fields
-  arr3Dvec ydownPhs; ///< Cache of phase shifts for calculating ydown fields
+  Flexible<Field2D> zShift;
+  Array<dcomplex> cmplx; ///< A temporary array, used for input/output to fft routines
+  Array<dcomplex> cmplxLoc; ///< A temporary array, used for input/output to fft routines
+
+  Matrix< Array<dcomplex> > getToAlignedPhs(CELL_LOC location = CELL_CENTRE); ///< Get phase shifts, calculating if necessary;
+  Matrix< Array<dcomplex> > getFromAlignedPhs(CELL_LOC location = CELL_CENTRE); ///< Get phase shifts, calculating if necessary;
+  Matrix< Array<dcomplex> > getYupPhs1(CELL_LOC location = CELL_CENTRE); ///< Get phase shifts, calculating if necessary;
+  Matrix< Array<dcomplex> > getYdownPhs1(CELL_LOC location = CELL_CENTRE); ///< Get phase shifts, calculating if necessary;
+  Matrix< Array<dcomplex> > getYupPhs2(CELL_LOC location = CELL_CENTRE); ///< Get phase shifts, calculating if necessary;
+  Matrix< Array<dcomplex> > getYdownPhs2(CELL_LOC location = CELL_CENTRE); ///< Get phase shifts, calculating if necessary;
+
+  bool has_toAligned_CENTRE, has_toAligned_XLOW, has_toAligned_YLOW; ///< Have phase shifts for shift to field aligned coordinates been calculated
+  bool has_fromAligned_CENTRE, has_fromAligned_XLOW, has_fromAligned_YLOW; ///< Have phase shifts for shift from field aligned coordinates been calculated
+  bool has_yupPhs1_CENTRE, has_yupPhs1_XLOW, has_yupPhs1_YLOW; ///< Have phase shifts for yup1 been calculated
+  bool has_ydownPhs1_CENTRE, has_ydownPhs1_XLOW, has_ydownPhs1_YLOW; ///< Have phase shifts for ydown1 been calculated
+  bool has_yupPhs2_CENTRE, has_yupPhs2_XLOW, has_yupPhs2_YLOW; ///< Have phase shifts for yup2 been calculated
+  bool has_ydownPhs2_CENTRE, has_ydownPhs2_XLOW, has_ydownPhs2_YLOW; ///< Have phase shifts for ydown2 been calculated
+
+  Matrix< Array<dcomplex> > toAlignedPhs_CENTRE; ///< Cache of phase shifts for transforming from X-Z orthogonal coordinates to field-aligned coordinates. Cell centre version.
+  Matrix< Array<dcomplex> > fromAlignedPhs_CENTRE; ///< Cache of phase shifts for transforming from field-aligned coordinates to X-Z orthogonal coordinates. Cell centre version.
+  Matrix< Array<dcomplex> > toAlignedPhs_XLOW; ///< Cache of phase shifts for transforming from X-Z orthogonal coordinates to field-aligned coordinates. Interpolated to CELL_XLOW.
+  Matrix< Array<dcomplex> > fromAlignedPhs_XLOW; ///< Cache of phase shifts for transforming from field-aligned coordinates to X-Z orthogonal coordinates. Interpolated to CELL_XLOW.
+  Matrix< Array<dcomplex> > toAlignedPhs_YLOW; ///< Cache of phase shifts for transforming from X-Z orthogonal coordinates to field-aligned coordinates. Interpolated to CELL_YLOW.
+  Matrix< Array<dcomplex> > fromAlignedPhs_YLOW; ///< Cache of phase shifts for transforming from field-aligned coordinates to X-Z orthogonal coordinates. Interpolated to CELL_YLOW.
+
+  Matrix< Array<dcomplex> > yupPhs1_CENTRE; ///< Cache of phase shifts for calculating yup1 fields. Cell centre version.
+  Matrix< Array<dcomplex> > ydownPhs1_CENTRE; ///< Cache of phase shifts for calculating ydown1 fields. Cell centre version.
+  Matrix< Array<dcomplex> > yupPhs2_CENTRE; ///< Cache of phase shifts for calculating yup2 fields. Cell centre version.
+  Matrix< Array<dcomplex> > ydownPhs2_CENTRE; ///< Cache of phase shifts for calculating ydown2 fields. Cell centre version.
+  Matrix< Array<dcomplex> > yupPhs1_XLOW; ///< Cache of phase shifts for calculating yup1 fields. Interpolated to CELL_XLOW.
+  Matrix< Array<dcomplex> > ydownPhs1_XLOW; ///< Cache of phase shifts for calculating ydown1 fields. Interpolated to CELL_XLOW.
+  Matrix< Array<dcomplex> > yupPhs2_XLOW; ///< Cache of phase shifts for calculating yup2 fields. Interpolated to CELL_XLOW.
+  Matrix< Array<dcomplex> > ydownPhs2_XLOW; ///< Cache of phase shifts for calculating ydown2 fields. Interpolated to CELL_XLOW.
+  Matrix< Array<dcomplex> > yupPhs1_YLOW; ///< Cache of phase shifts for calculating yup1 fields. Interpolated to CELL_YLOW.
+  Matrix< Array<dcomplex> > ydownPhs1_YLOW; ///< Cache of phase shifts for calculating ydown1 fields. Interpolated to CELL_YLOW.
+  Matrix< Array<dcomplex> > yupPhs2_YLOW; ///< Cache of phase shifts for calculating yup2 fields. Interpolated to CELL_YLOW.
+  Matrix< Array<dcomplex> > ydownPhs2_YLOW; ///< Cache of phase shifts for calculating ydown2 fields. Interpolated to CELL_YLOW.
 
   /*!
    * Shift a 2D field in Z. 
    * Since 2D fields are constant in Z, this has no effect
    */
-  const Field2D shiftZ(const Field2D &f, const Field2D &UNUSED(zangle)){return f;};
+  const Field2D shiftZ(const Field2D &f, const Field2D &UNUSED(zangle), const REGION UNUSED(region)=RGN_NOX){return f;};
 
   /*!
    * Shift a 3D field \p f in Z by the given \p zangle
@@ -152,7 +187,7 @@ private:
    * @param[in] zangle   Toroidal angle (z)
    *
    */ 
-  const Field3D shiftZ(const Field3D &f, const Field2D &zangle);
+  const Field3D shiftZ(const Field3D &f, const Field2D &zangle, const REGION region=RGN_NOX);
 
   /*!
    * Shift a 3D field \p f by the given phase \p phs in Z
@@ -163,7 +198,7 @@ private:
    * @param[in] f  The field to shift
    * @param[in] phs  The phase to shift by
    */
-  const Field3D shiftZ(const Field3D &f, const arr3Dvec &phs);
+  const Field3D shiftZ(const Field3D &f, const Matrix< Array<dcomplex> > &phs, const REGION region=RGN_NOX);
 
   /*!
    * Shift a given 1D array, assumed to be in Z, by the given \p zangle
@@ -182,7 +217,10 @@ private:
    * @param[in] phs Phase shift, assumed to have length (mesh.LocalNz/2 + 1) i.e. the number of modes
    * @param[out] out  A 1D array of length mesh.LocalNz, already allocated
    */
-  void shiftZ(const BoutReal *in, const std::vector<dcomplex> &phs, BoutReal *out);
+  void shiftZ(const BoutReal *in, const Array<dcomplex> &phs, BoutReal *out);
+
+  /// Write out ParallelTransform variables to file
+  void outputVars(Datafile &file);
 };
 
 

include/field3d.hxx

@@ -151,21 +151,21 @@ template <typename F> class Flexible;
 
       Field3D f(0.0); // f allocated, set to zero
 
-      f.yup() // error; f.yup not allocated
+      f.yup(1) // error; f.yup not allocated
 
-      f.mergeYupYdown(); // f.yup() and f.ydown() now point to f
-      f.yup()(0,1,0)  // ok, gives value of f at (0,1,0)
+      f.mergeYupYdown(); // f.yup(i) and f.ydown(i) now point to f
+      f.yup(1)(0,1,0)  // ok, gives value of f at (0,1,0)
 
       To have separate fields for yup and ydown, first call
 
-      f.splitYupYdown(); // f.yup() and f.ydown() separate
+      f.splitYupYdown(); // f.yup(i) and f.ydown(i) separate
 
-      f.yup(); // ok
-      f.yup()(0,1,0) // error; f.yup not allocated
+      f.yup(1); // ok
+      f.yup(1)(0,1,0) // error; f.yup not allocated
 
-      f.yup() = 1.0; // Set f.yup() field to 1.0
+      f.yup(1) = 1.0; // Set f.yup() field to 1.0
 
-      f.yup()(0,1,0) // ok
+      f.yup(1)(0,1,0) // ok
 
  */
 class Field3D : public Field, public FieldData {
@@ -234,36 +234,30 @@ class Field3D : public Field, public FieldData {
    * Ensure that yup and ydown refer to this field
    */
   void mergeYupYdown();
+
+  /*!
+   * Clear all yup/ydown fields and field_fa
+   */
+  void clearYupYdown();
 
   /// Check if this field has yup and ydown fields
   bool hasYupYdown() const {
-    return (yup_field != nullptr) && (ydown_field != nullptr);
+    return (yup1_field != nullptr) && (ydown1_field != nullptr);
   }
 
   /// Return reference to yup field
-  Field3D& yup() { 
-    ASSERT2(yup_field != nullptr); // Check for communicate
-    return *yup_field; 
-  }
-  /// Return const reference to yup field
-  const Field3D& yup() const { 
-    ASSERT2(yup_field != nullptr);
-    return *yup_field; 
-  }
+  Field3D& yup(const int i = 1);
 
+  /// Return const reference to yup field
+  const Field3D& yup(const int i = 1) const;
+
   /// Return reference to ydown field
-  Field3D& ydown() { 
-    ASSERT2(ydown_field != nullptr);
-    return *ydown_field;
-  }
+  Field3D& ydown(const int i = 1);
 
   /// Return const reference to ydown field
-  const Field3D& ydown() const { 
-    ASSERT2(ydown_field != nullptr);
-    return *ydown_field; 
-  }
+  const Field3D& ydown(const int i = 1) const;
 
-  /// Return yup if dir=+1, and ydown if dir=-1
+  /// Return yup(dir) if dir>0, and ydown(-dir) if dir<0
   Field3D& ynext(int dir);
   const Field3D& ynext(int dir) const;
 
@@ -528,8 +522,8 @@ private:
 
   Field3D *deriv; ///< Time derivative (may be NULL)
 
-  /// Pointers to fields containing values along Y
-  Field3D *yup_field, *ydown_field;
+  /// Arrays of pointers to fields containing values along Y
+  Field3D *yup1_field, *ydown1_field, *yup2_field, *ydown2_field;
 };
 
 // Non-member overloaded operators