New shifted metric

include/bout/mesh.hxx

@@ -609,6 +609,12 @@ class Mesh {
   /// Derivative functions of a velocity field, and field stencil v, f
   typedef BoutReal (*flux_func)(stencil&, stencil &);
 
+  /// Calculate yup/ydown fields in case they can be computed for an
+  /// intermediate variable without communicating
+  void calcYUpDown(Field3D &f, REGION region = RGN_NOX) {
+    getParallelTransform().calcYUpDown(f, region);
+  }
+
   /// Transform a field into field-aligned coordinates
   const Field3D toFieldAligned(const Field3D &f, const REGION region = RGN_NOX) {
     return getParallelTransform().toFieldAligned(f, region);

include/bout/paralleltransform.hxx

@@ -10,6 +10,7 @@
 #include <boutexception.hxx>
 #include <dcomplex.hxx>
 #include <unused.hxx>
+#include <utils.hxx>
 
 class Mesh;
 
@@ -28,7 +29,7 @@ public:
   virtual ~ParallelTransform() {}
 
   /// Given a 3D field, calculate and set the Y up down fields
-  virtual void calcYUpDown(Field3D &f) = 0;
+  virtual void calcYUpDown(Field3D &f, REGION region = RGN_NOX) = 0;
 
   /// Calculate Yup and Ydown fields by integrating over mapped points
   /// This should be used for parallel divergence operators
@@ -59,7 +60,7 @@ public:
    * Merges the yup and ydown() fields of f, so that
    * f.yup() = f.ydown() = f
    */ 
-  void calcYUpDown(Field3D &f) override {f.mergeYupYdown();}
+  void calcYUpDown(Field3D &f, REGION UNUSED(region) = RGN_NOX) override {f.mergeYupYdown();}
 
   /*!
    * The field is already aligned in Y, so this
@@ -99,7 +100,7 @@ public:
    * Calculates the yup() and ydown() fields of f
    * by taking FFTs in Z and applying a phase shift.
    */ 
-  void calcYUpDown(Field3D &f) override;
+  void calcYUpDown(Field3D &f, REGION UNUSED(region) = RGN_NOX) override;
 
   /*!
    * Uses FFTs and a phase shift to align the grid points
@@ -184,5 +185,115 @@ private:
   void shiftZ(const BoutReal *in, const std::vector<dcomplex> &phs, BoutReal *out);
 };
 
+/*!
+ * Alternative shifted metric method
+ * Fields are stored on a grid is orthogonal in X-Z, interpolated onto a
+ * field-aligned grid to calculate parallel derivatives or interpolations
+ *
+ * In this implementation the interpolation is done using FFTs in Z
+ */
+class ShiftToFieldAligned : public ParallelTransform {
+public:
+  ShiftToFieldAligned(Mesh &mesh);
+
+  /*!
+   * Do not calculate yup() and ydown() fields of f
+   * for this method.
+   * Instead use this method to calculate f.field_fa.
+   * This is a bit hacky, but we can rename ParallelTransform::calcYUpDown to
+   * something more generic later.
+   */
+  void calcYUpDown(Field3D &f, REGION region = RGN_NOX) override;
+
+  /*!
+   * Uses FFTs and a phase shift to align the grid points
+   * with the y coordinate (along magnetic field usually).
+   *
+   * Note that the returned field will no longer be orthogonal
+   * in X-Z, and the metric tensor will need to be changed
+   * if X derivatives are used.
+   */
+  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, const REGION region=RGN_NOX) override;
+
+  bool canToFromFieldAligned() override{
+    return true;
+  }
+
+private:
+  ShiftToFieldAligned();
+
+  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_CENTRE, zShift_XLOW, zShift_YLOW; ///< The angles to shift from orthogonal to field-aligned coordinates
+  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;
+
+  bool has_toAligned_CENTRE, has_toAligned_XLOW, has_toAligned_YLOW; ///< Flags saying whether phases for shifts to field-aligned coordinates have been calculated.
+  bool has_fromAligned_CENTRE, has_fromAligned_XLOW, has_fromAligned_YLOW; ///< Flags saying whether phases for shifts from field-aligned coordinates have 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.
+
+  /*!
+   * 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), const REGION UNUSED(region)=RGN_NOX){return f;};
+
+  /*!
+   * Shift a 3D field \p f in Z by the given \p zangle
+   *
+   * @param[in] f  The field to shift
+   * @param[in] zangle   Toroidal angle (z)
+   *
+   */
+  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
+   *
+   * Calculates FFT in Z, multiplies by the complex phase
+   * and inverse FFTS.
+   *
+   * @param[in] f  The field to shift
+   * @param[in] phs  The phase to shift by
+   */
+  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
+   *
+   * @param[in] in  A 1D array of length \p len
+   * @param[in] len  Length of the in and out arrays
+   * @param[in] zangle  The angle (z coordinate) to shift by
+   * @param[out] out  A 1D array of length \p len, already allocated
+   */
+  void shiftZ(const BoutReal *in, int len, BoutReal zangle,  BoutReal *out);
+
+  /*!
+   * Shift a given 1D array, assumed to be in Z, by the given \p zangle
+   *
+   * @param[in] in  A 1D array of length mesh.LocalNz
+   * @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 Array<dcomplex> &phs, BoutReal *out);
+};
+
 
 #endif // __PARALLELTRANSFORM_H__

include/field3d.hxx

@@ -154,6 +154,11 @@ class Mesh;  // #include "bout/mesh.hxx"
 
       f.yup()(0,1,0) // ok
 
+  It is also useful to be able to access a version of the field interpolated
+  (in the z-direction) onto a field-aligned grid. This can be
+  calculated/accessed with the Mesh::toFieldAligned() and the main field can be
+  interpolated from a field-aligned one with Mesh::fromFieldAligned().
+
  */
 class Field3D : public Field, public FieldData {
  public:
@@ -234,6 +239,19 @@ class Field3D : public Field, public FieldData {
     return (yup_field != nullptr) && (ydown_field != nullptr);
   }
 
+  /// Check if this field has a field-aligned version
+  bool hasFieldAligned() const {
+    return has_field_aligned;
+  }
+
+  /// Set has_field_aligned
+  void setHasFieldAligned(bool val) {
+    if (val) {
+      ASSERT1(field_fa != nullptr && field_fa->isAllocated());
+    }
+    has_field_aligned = val;
+  }
+
   /// Return reference to yup field
   Field3D& yup() { 
     ASSERT2(yup_field != nullptr); // Check for communicate
@@ -261,6 +279,12 @@ class Field3D : public Field, public FieldData {
   Field3D& ynext(int dir);
   const Field3D& ynext(int dir) const;
 
+  /// Return reference to field-aligned field
+  Field3D& fieldAligned();
+
+  /// Return const reference to field-aligned field
+  const Field3D& fieldAligned() const;
+
   /// Set variable location for staggered grids to @param new_location
   ///
   /// Throws BoutException if new_location is not `CELL_CENTRE` and
@@ -464,6 +488,10 @@ private:
 
   /// Pointers to fields containing values along Y
   Field3D *yup_field, *ydown_field;
+
+  /// Pointer to field containing the field-aligned version of the field
+  Field3D *field_fa;
+  bool has_field_aligned;
 };
 
 // Non-member overloaded operators

src/field/field3d.cxx

@@ -46,7 +46,8 @@
 /// Constructor
 Field3D::Field3D(Mesh *localmesh)
     : Field(localmesh), background(nullptr), deriv(nullptr), yup_field(nullptr),
-      ydown_field(nullptr) {
+      ydown_field(nullptr),
+      field_fa(nullptr), has_field_aligned(false) {
 #ifdef TRACK
   name = "<F3D>";
 #endif
@@ -72,7 +73,8 @@ Field3D::Field3D(Mesh *localmesh)
 Field3D::Field3D(const Field3D &f)
     : Field(f.fieldmesh),                // The mesh containing array sizes
       background(nullptr), data(f.data), // This handles references to the data array
-      deriv(nullptr), yup_field(nullptr), ydown_field(nullptr) {
+      deriv(nullptr), yup_field(nullptr), ydown_field(nullptr),
+      field_fa(nullptr), has_field_aligned(false) {
 
   TRACE("Field3D(Field3D&)");
 
@@ -101,7 +103,8 @@ Field3D::Field3D(const Field3D &f)
 
 Field3D::Field3D(const Field2D &f)
     : Field(f.getMesh()), background(nullptr), deriv(nullptr), yup_field(nullptr),
-      ydown_field(nullptr) {
+      ydown_field(nullptr),
+      field_fa(nullptr), has_field_aligned(false) {
 
   TRACE("Field3D: Copy constructor from Field2D");
 
@@ -119,7 +122,8 @@ Field3D::Field3D(const Field2D &f)
 
 Field3D::Field3D(const BoutReal val, Mesh *localmesh)
     : Field(localmesh), background(nullptr), deriv(nullptr), yup_field(nullptr),
-      ydown_field(nullptr) {
+      ydown_field(nullptr),
+      field_fa(nullptr), has_field_aligned(false) {
 
   TRACE("Field3D: Copy constructor from value");
 
@@ -199,15 +203,20 @@ void Field3D::mergeYupYdown() {
 
 void Field3D::deleteYupYdown() {
   // Delete auxiliary fields if they have been set
-  if (yup_field == this && ydown_field == this) {
-    return;
-  }
+  if (!(yup_field == this && ydown_field == this)) {
+
+    delete yup_field;
+    yup_field = nullptr;
 
-  delete yup_field;
-  yup_field = nullptr;
+    delete ydown_field;
+    ydown_field = nullptr;
+  }
 
-  delete ydown_field;
-  ydown_field = nullptr;
+  if (field_fa != this) {
+    delete field_fa;
+  }
+  field_fa = nullptr;
+  has_field_aligned = false;
 }
 
 Field3D& Field3D::ynext(int dir) {
@@ -232,6 +241,25 @@ const Field3D& Field3D::ynext(int dir) const {
   }
 }
 
+
+Field3D& Field3D::fieldAligned() {
+  if (field_fa == nullptr) {
+    // has_field_aligned should be false now, until field_fa is set to a value.
+    // This should have been set either in the constructor or when field_fa was
+    // deleted.
+    ASSERT1(!has_field_aligned);
+
+    field_fa = new Field3D(getMesh());
+    field_fa->setLocation(location);
+  }
+  return *field_fa;
+}
+
+const Field3D& Field3D::fieldAligned() const {
+  ASSERT1(field_fa != nullptr);
+  return *field_fa;
+}
+
 void Field3D::setLocation(CELL_LOC new_location) {
   if (getMesh()->StaggerGrids) {
     if (new_location == CELL_VSHIFT) {

src/mesh/coordinates.cxx

@@ -702,26 +702,30 @@ const Field3D Coordinates::Div_par(const Field3D &f, CELL_LOC outloc,
 
   // Need Bxy at location of f, which might be different from location of this
   // Coordinates object
-  Field2D Bxy_floc = f.getCoordinates()->Bxy;
-
-  if (!f.hasYupYdown()) {
-    // No yup/ydown fields. The Grad_par operator will
-    // shift to field aligned coordinates
-    return Bxy * Grad_par(f / Bxy_floc, outloc, method);
-  }
-
-  // Need to modify yup and ydown fields
+  Field2D& Bxy_floc = f.getCoordinates()->Bxy;
   Field3D f_B = f / Bxy_floc;
-  if (&f.yup() == &f) {
-    // Identity, yup and ydown point to same field
-    f_B.mergeYupYdown();
+
+  if (f.hasYupYdown()) {
+    // Need to modify yup and ydown fields
+    if (&f.yup() == &f) {
+      // Identity, yup and ydown point to same field
+      f_B.mergeYupYdown();
+    } else {
+      // Distinct fields
+      f_B.splitYupYdown();
+      f_B.yup() = f.yup() / Bxy_floc;
+      f_B.ydown() = f.ydown() / Bxy_floc;
+    }
+    return Bxy * Grad_par(f_B, outloc, method);
+  } else if (f.hasFieldAligned()) {
+    f_B.fieldAligned() = f.fieldAligned() / Bxy;
+    f_B.setHasFieldAligned(true);
+    return Bxy * Grad_par(f_B, outloc, method);
   } else {
-    // Distinct fields
-    f_B.splitYupYdown();
-    f_B.yup() = f.yup() / Bxy_floc;
-    f_B.ydown() = f.ydown() / Bxy_floc;
+    // No yup/ydown or field-aligned fields. The Grad_par operator will shift
+    // to field aligned coordinates
+    return Bxy * Grad_par(f_B, outloc, method);
   }
-  return Bxy * Grad_par(f_B, outloc, method);
 }
 
 /////////////////////////////////////////////////////////

src/mesh/mesh.cxx

@@ -292,23 +292,27 @@ void Mesh::setParallelTransform() {
 
   // Convert to lower case for comparison
   ptstr = lowercase(ptstr);
-    
+
   if(ptstr == "identity") {
     // Identity method i.e. no transform needed
     transform = std::unique_ptr<ParallelTransform>(new ParallelTransformIdentity());
-      
+
   }else if(ptstr == "shifted") {
     // Shifted metric method
     transform = std::unique_ptr<ParallelTransform>(new ShiftedMetric(*this));
-
+
+  }else if(ptstr == "shifttofieldaligned") {
+    // Alternative shifted metric method
+    transform = std::unique_ptr<ParallelTransform>(new ShiftToFieldAligned(*this));
+
   }else if(ptstr == "fci") {
 
     Options *fci_options = Options::getRoot()->getSection("fci");
     // Flux Coordinate Independent method
     bool fci_zperiodic;
     fci_options->get("z_periodic", fci_zperiodic, true);
     transform = std::unique_ptr<ParallelTransform>(new FCITransform(*this, fci_zperiodic));
-      
+
   }else {
     throw BoutException("Unrecognised paralleltransform option.\n"
                         "Valid choices are 'identity', 'shifted', 'fci'");